-
Notifications
You must be signed in to change notification settings - Fork 129
/
proxying_reports_server.go
47 lines (39 loc) · 1.66 KB
/
proxying_reports_server.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
package scheduler
import (
"context"
"time"
"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
)
type ProxyingSchedulingReportsServer struct {
client schedulerobjects.SchedulerReportingClient
}
func NewProxyingSchedulingReportsServer(client schedulerobjects.SchedulerReportingClient) *ProxyingSchedulingReportsServer {
return &ProxyingSchedulingReportsServer{
client: client,
}
}
func (s *ProxyingSchedulingReportsServer) GetSchedulingReport(ctx context.Context, request *schedulerobjects.SchedulingReportRequest) (*schedulerobjects.SchedulingReport, error) {
ctx, cancel := reduceTimeout(ctx)
defer cancel()
return s.client.GetSchedulingReport(ctx, request)
}
func (s *ProxyingSchedulingReportsServer) GetQueueReport(ctx context.Context, request *schedulerobjects.QueueReportRequest) (*schedulerobjects.QueueReport, error) {
ctx, cancel := reduceTimeout(ctx)
defer cancel()
return s.client.GetQueueReport(ctx, request)
}
func (s *ProxyingSchedulingReportsServer) GetJobReport(ctx context.Context, request *schedulerobjects.JobReportRequest) (*schedulerobjects.JobReport, error) {
ctx, cancel := reduceTimeout(ctx)
defer cancel()
return s.client.GetJobReport(ctx, request)
}
// We reduce the context deadline here, to prevent our call and the caller who called us from timing out at the same time
// This should mean our caller gets the real error message rather than a generic timeout error from client side
func reduceTimeout(ctx context.Context) (context.Context, context.CancelFunc) {
deadline, ok := ctx.Deadline()
if !ok {
return context.WithCancel(ctx)
}
deadline = deadline.Add(-time.Second)
return context.WithDeadline(ctx, deadline)
}