Merge pull request #38 from lrascao/backport-get-work-items-send-timeout

Merge pull request #24 from famarting/get-work-items-send-timeout

Author: cicoyle

backend/executor.go

@@ -56,6 +56,7 @@ type grpcExecutor struct {
 	onWorkItemConnection func(context.Context) error
 	onWorkItemDisconnect func(context.Context) error
 	streamShutdownChan   <-chan any
+	streamSendTimeout    *time.Duration
 }
 
 type grpcExecutorOptions func(g *grpcExecutor)
@@ -90,6 +91,12 @@ func WithStreamShutdownChannel(c <-chan any) grpcExecutorOptions {
 	}
 }
 
+func WithStreamSendTimeout(d time.Duration) grpcExecutorOptions {
+	return func(g *grpcExecutor) {
+		g.streamSendTimeout = &d
+	}
+}
+
 // NewGrpcExecutor returns the Executor object and a method to invoke to register the gRPC server in the executor.
 func NewGrpcExecutor(be Backend, logger Logger, opts ...grpcExecutorOptions) (executor Executor, registerServerFn func(grpcServer grpc.ServiceRegistrar)) {
 	grpcExecutor := &grpcExecutor{
@@ -322,7 +329,7 @@ func (g *grpcExecutor) GetWorkItems(req *protos.GetWorkItemsRequest, stream prot
 				}
 			}
 
-			if err := stream.Send(wi); err != nil {
+			if err := g.sendWorkItem(stream, wi); err != nil {
 				g.logger.Errorf("encountered an error while sending work item: %v", err)
 				return err
 			}
@@ -336,6 +343,27 @@ func (g *grpcExecutor) GetWorkItems(req *protos.GetWorkItemsRequest, stream prot
 	}
 }
 
+func (g *grpcExecutor) sendWorkItem(stream protos.TaskHubSidecarService_GetWorkItemsServer, wi *protos.WorkItem) error {
+	ctx := stream.Context()
+	if g.streamSendTimeout != nil {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, *g.streamSendTimeout)
+		defer cancel()
+	}
+
+	errCh := make(chan error, 2)
+	go func() {
+		select {
+		case errCh <- stream.Send(wi):
+		case <-ctx.Done():
+			g.logger.Errorf("timed out while sending work item")
+			errCh <- ctx.Err()
+		}
+	}()
+
+	return <-errCh
+}
+
 // CompleteOrchestratorTask implements protos.TaskHubSidecarServiceServer
 func (g *grpcExecutor) CompleteOrchestratorTask(ctx context.Context, res *protos.OrchestratorResponse) (*protos.CompleteTaskResponse, error) {
 	iid := api.InstanceID(res.InstanceId)