Repository: numaproj/numaflow
Pull request #1515 (2024-02-20 05:36:58)

fix: avoid panic when ctx is canceled


Patch:
diff --git a/pkg/reduce/pnf/processandforward.go b/pkg/reduce/pnf/processandforward.go
index c2be8e83..7f793510 100644
--- a/pkg/reduce/pnf/processandforward.go
+++ b/pkg/reduce/pnf/processandforward.go
@@ -133,7 +133,7 @@ outerLoop:
 	for {
 		select {
 		case err := <-errCh:
-			if errors.Is(err, ctx.Err()) {
+			if errors.Is(err, context.Canceled) {
 				return
 			}
 			if err != nil {
@@ -170,7 +170,7 @@ outerLoop:
 	for {
 		select {
 		case err := <-errCh:
-			if errors.Is(err, ctx.Err()) {
+			if errors.Is(err, context.Canceled) {
 				return
 			}
 			if err != nil {
diff --git a/pkg/sdkclient/error/error.go b/pkg/sdkclient/error/error.go
index 7c4ffbfb..5f58fc72 100644
--- a/pkg/sdkclient/error/error.go
+++ b/pkg/sdkclient/error/error.go
@@ -26,6 +26,7 @@ type ErrKind int16
 const (
 	Retryable    ErrKind = iota // The error is retryable
 	NonRetryable                // The error is non-retryable
+	Canceled                    // Request canceled
 	Unknown                     // Unknown err kind
 )
 
@@ -35,6 +36,8 @@ func (ek ErrKind) String() string {
 		return "Retryable"
 	case NonRetryable:
 		return "NonRetryable"
+	case Canceled:
+		return "Canceled"
 	case Unknown:
 		return "Unknown"
 	default:
diff --git a/pkg/shared/util/grpc_utils.go b/pkg/shared/util/grpc_utils.go
index 8c453f61..0f189e0a 100644
--- a/pkg/shared/util/grpc_utils.go
+++ b/pkg/shared/util/grpc_utils.go
@@ -55,6 +55,9 @@ func ToUDFErr(name string, err error) error {
 		udfError = sdkerr.New(sdkerr.Retryable, statusCode.Message())
 		log.Printf("failed %s: %s", name, udfError.Error())
 		return udfError
+	case codes.Canceled:
+		udfError = sdkerr.New(sdkerr.Canceled, statusCode.Message())
+		return udfError
 	default:
 		log.Printf("failed %s: %s", name, udfError.Error())
 		return udfError
diff --git a/pkg/udf/rpc/errors.go b/pkg/udf/rpc/errors.go
index 083e45c5..c711dcd0 100644
--- a/pkg/udf/rpc/errors.go
+++ b/pkg/udf/rpc/errors.go
@@ -44,3 +44,8 @@ func (e ApplyUDFErr) IsInternalErr() bool {
 func (e ApplyUDFErr) Error() string {
 	return fmt.Sprint(e.Message)
 }
+
+// Is checks if the error is of the same type
+func (e ApplyUDFErr) Is(target error) bool {
+	return target.Error() == e.Error()
+}
diff --git a/pkg/udf/rpc/grpc_aligned_reduce.go b/pkg/udf/rpc/grpc_aligned_reduce.go
index d4522dbf..c26aec2c 100644
--- a/pkg/udf/rpc/grpc_aligned_reduce.go
+++ b/pkg/udf/rpc/grpc_aligned_reduce.go
@@ -227,6 +227,15 @@ func convertToUdfError(err error) ApplyUDFErr {
 				MainCarDown: false,
 			},
 		}
+	case sdkerr.Canceled:
+		return ApplyUDFErr{
+			UserUDFErr: false,
+			Message:    context.Canceled.Error(),
+			InternalErr: InternalErr{
+				Flag:        true,
+				MainCarDown: false,
+			},
+		}
 	default:
 		return ApplyUDFErr{
 			UserUDFErr: false,


Test patch:
diff --git a/pkg/sdkclient/error/error_test.go b/pkg/sdkclient/error/error_test.go
index 4155bf8a..b973b03f 100644
--- a/pkg/sdkclient/error/error_test.go
+++ b/pkg/sdkclient/error/error_test.go
@@ -39,9 +39,9 @@ func TestErrKind_String(t *testing.T) {
 			want: "NonRetryable",
 		},
 		{
-			name: "unknown_1",
+			name: "canceled",
 			ek:   2,
-			want: "Unknown",
+			want: "Canceled",
 		},
 		{
 			name: "unknown_2",
@@ -94,7 +94,7 @@ func TestFromError(t *testing.T) {
 			name: "good_unknown_1",
 			args: args{
 				err: &UDFError{
-					errKind:    2,
+					errKind:    3,
 					errMessage: "Unknown Error",
 				},
 			},
@@ -109,6 +109,17 @@ func TestFromError(t *testing.T) {
 			want:   &UDFError{Unknown, "not a standard error"},
 			wantOk: false,
 		},
+		{
+			name: "good_canceled",
+			args: args{
+				err: &UDFError{
+					errKind:    2,
+					errMessage: "context canceled",
+				},
+			},
+			want:   &UDFError{Canceled, "context canceled"},
+			wantOk: true,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {