Handle cancellation of GRPC stream. Closes #17

Author: kaidaguerre

grpc/shared/grpc.go

@@ -20,8 +20,10 @@ func (c *GRPCClient) GetSchema(req *proto.GetSchemaRequest) (*proto.GetSchemaRes
 	return c.client.GetSchema(c.ctx, req)
 }
 
-func (c *GRPCClient) Execute(req *proto.ExecuteRequest) (proto.WrapperPlugin_ExecuteClient, error) {
-	return c.client.Execute(c.ctx, req)
+func (c *GRPCClient) Execute(req *proto.ExecuteRequest) (proto.WrapperPlugin_ExecuteClient, context.CancelFunc, error) {
+	ctx, cancel := context.WithCancel(c.ctx)
+	client, err := c.client.Execute(ctx, req)
+	return client, cancel, err
 }
 
 func (c *GRPCClient) SetConnectionConfig(req *proto.SetConnectionConfigRequest) (*proto.SetConnectionConfigResponse, error) {

grpc/shared/interface.go

@@ -29,7 +29,7 @@ type WrapperPluginServer interface {
 
 type WrapperPluginClient interface {
 	GetSchema(request *proto.GetSchemaRequest) (*proto.GetSchemaResponse, error)
-	Execute(req *proto.ExecuteRequest) (proto.WrapperPlugin_ExecuteClient, error)
+	Execute(req *proto.ExecuteRequest) (proto.WrapperPlugin_ExecuteClient, context.CancelFunc, error)
 	SetConnectionConfig(req *proto.SetConnectionConfigRequest) (*proto.SetConnectionConfigResponse, error)
 }
 

plugin/errors.go

@@ -0,0 +1,3 @@
+package plugin
+
+const contextCancelledError = "rpc error: code = Canceled desc = context canceled"

plugin/hydrate.go

@@ -36,8 +36,13 @@ func WrapHydrate(hydrateFunc HydrateFunc, shouldIgnoreError ErrorPredicate) Hydr
 	return func(ctx context.Context, d *QueryData, h *HydrateData) (item interface{}, err error) {
 		defer func() {
 			if r := recover(); r != nil {
-				log.Printf("[WARN] recovered a panic from a wrapped hydrate function: %v\n", r)
-				err = status.Error(codes.Internal, fmt.Sprintf("hydrate function %s failed with panic %v", helpers.GetFunctionName(hydrateFunc), r))
+				if helpers.ToError(r).Error() == contextCancelledError {
+					// if the error was a context cancellation, just trace it - this is not an error
+					log.Printf("[TRACE] hydrate function %s terminated with a context cancellation", helpers.GetFunctionName(hydrateFunc))
+				} else {
+					log.Printf("[WARN] recovered a panic from a wrapped hydrate function: %v\n", r)
+					err = status.Error(codes.Internal, fmt.Sprintf("hydrate function %s failed with panic %v", helpers.GetFunctionName(hydrateFunc), r))
+				}
 			}
 		}()
 		// call the underlying get function

plugin/query_data.go

@@ -2,6 +2,7 @@ package plugin
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log"
 	"sync"
@@ -54,6 +55,9 @@ type QueryData struct {
 	listWg sync.WaitGroup
 	// when executing parent child list calls, we cache the parent list result in the query data passed to the child list call
 	parentItem interface{}
+
+	// there was an error streaming to the grpc stream
+	streamingError error
 }
 
 func newQueryData(queryContext *QueryContext, table *Table, stream proto.WrapperPlugin_ExecuteServer, connection *Connection, matrix []map[string]interface{}, connectionManager *connection_manager.Manager) *QueryData {
@@ -261,6 +265,11 @@ func (d *QueryData) verifyCallerIsListCall(callingFunction string) bool {
 }
 
 func (d *QueryData) streamLeafListItem(ctx context.Context, item interface{}) {
+	if d.streamingError != nil {
+		// if there is streaming error, panic to force exit thread - this will be recovered higher up
+		panic(d.streamingError)
+	}
+
 	// create rowData, passing matrixItem from context
 	rd := newRowData(d, item)
 	rd.matrixItem = GetMatrixItem(ctx)
@@ -283,6 +292,8 @@ func (d *QueryData) streamRows(_ context.Context, rowChan chan *proto.Row) error
 	for {
 		// wait for either an item or an error
 		select {
+		case <-d.stream.Context().Done():
+			d.streamingError = errors.New(contextCancelledError)
 		case err := <-d.errorChan:
 			log.Printf("[ERROR] streamRows error chan select: %v\n", err)
 			return err
@@ -295,7 +306,10 @@ func (d *QueryData) streamRows(_ context.Context, rowChan chan *proto.Row) error
 				return nil
 			}
 			if err := d.streamRow(row); err != nil {
-				log.Printf("[ERROR] stream.Send returned error: %v\n", err)
+				// if there was an error streaming, store in d.streamingError
+				// - this is checked by the thread streaming list items and will cause it to terminate
+				d.streamingError = err
+
 				return err
 			}
 		}