Fix handling of context errors

common/persistence/visibility/store/elasticsearch/visibility_store.go

@@ -440,7 +440,7 @@ func (s *VisibilityStore) countGroupByWorkflowExecutions(
 		termsAgg,
 	)
 	if err != nil {
-		return nil, err
+		return nil, ConvertElasticsearchClientError("CountWorkflowExecutions failed", err)
 	}
 	return s.parseCountGroupByResponse(esResponse, groupByFields)
 }
@@ -1154,14 +1154,19 @@ func finishParseJSONValue(val interface{}, t enumspb.IndexedValueType) (interfac
 
 func ConvertElasticsearchClientError(message string, err error) error {
 	errMessage := fmt.Sprintf("%s: %s", message, detailedErrorMessage(err))
-	switch e := err.(type) {
-	case *elastic.Error:
-		switch e.Status {
+	var elasticErr *elastic.Error
+	switch {
+	case errors.As(err, &elasticErr):
+		switch elasticErr.Status {
 		case 400: // BadRequest
 			// Returning InvalidArgument error will prevent retry on a caller side.
 			return serviceerror.NewInvalidArgument(errMessage)
 		}
+		return serviceerror.NewUnavailable(errMessage)
+	case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded):
+		return fmt.Errorf("%s: %w", message, err)
 	}
+
 	return serviceerror.NewUnavailable(errMessage)
 }
 

common/persistence/visibility/store/sql/visibility_store.go

@@ -73,6 +73,13 @@ func (s *VisibilityStore) GetName() string {
 	return s.sqlStore.GetName()
 }
 
+func convertSQLError(message string, err error) error {
+	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+		return fmt.Errorf("%s: %w", message, err)
+	}
+	return serviceerror.NewUnavailable(fmt.Sprintf("%s: %v", message, err))
+}
+
 func (s *VisibilityStore) GetIndexName() string {
 	return s.sqlStore.GetDbName()
 }
@@ -160,7 +167,7 @@ func (s *VisibilityStore) DeleteWorkflowExecution(
 		RunID:       request.RunID,
 	})
 	if err != nil {
-		return serviceerror.NewUnavailable(err.Error())
+		return convertSQLError("DeleteWorkflowExecution operation failed.", err)
 	}
 	return nil
 }
@@ -225,8 +232,7 @@ func (s *VisibilityStore) listWorkflowExecutions(
 
 	rows, err := s.sqlStore.DB.SelectFromVisibility(ctx, *selectFilter)
 	if err != nil {
-		return nil, serviceerror.NewUnavailable(
-			fmt.Sprintf("ListWorkflowExecutions operation failed. Select failed: %v", err))
+		return nil, convertSQLError("ListWorkflowExecutions operation failed.", err)
 	}
 	if len(rows) == 0 {
 		return &store.InternalListWorkflowExecutionsResponse{}, nil
@@ -296,8 +302,7 @@ func (s *VisibilityStore) listWorkflowExecutionsLegacy(
 
 	rows, err := s.sqlStore.DB.SelectFromVisibility(ctx, *selectFilter)
 	if err != nil {
-		return nil, serviceerror.NewUnavailable(
-			fmt.Sprintf("ListWorkflowExecutions operation failed. Select failed: %v", err))
+		return nil, convertSQLError("ListWorkflowExecutions operation failed.", err)
 	}
 	if len(rows) == 0 {
 		return &store.InternalListWorkflowExecutionsResponse{}, nil
@@ -381,8 +386,7 @@ func (s *VisibilityStore) countWorkflowExecutionsLegacy(
 
 	count, err := s.sqlStore.DB.CountFromVisibility(ctx, *selectFilter)
 	if err != nil {
-		return nil, serviceerror.NewUnavailable(
-			fmt.Sprintf("CountWorkflowExecutions operation failed. Query failed: %v", err))
+		return nil, convertSQLError("CountWorkflowExecutions operation failed.", err)
 	}
 
 	return &manager.CountWorkflowExecutionsResponse{Count: count}, nil
@@ -443,8 +447,7 @@ func (s *VisibilityStore) countWorkflowExecutions(
 
 	count, err := s.sqlStore.DB.CountFromVisibility(ctx, *selectFilter)
 	if err != nil {
-		return nil, serviceerror.NewUnavailable(
-			fmt.Sprintf("CountWorkflowExecutions operation failed. Query failed: %v", err))
+		return nil, convertSQLError("CountWorkflowExecutions operation failed.", err)
 	}
 
 	return &manager.CountWorkflowExecutionsResponse{Count: count}, nil
@@ -469,8 +472,7 @@ func (s *VisibilityStore) countGroupByWorkflowExecutions(
 
 	rows, err := s.sqlStore.DB.CountGroupByFromVisibility(ctx, *selectFilter)
 	if err != nil {
-		return nil, serviceerror.NewUnavailable(
-			fmt.Sprintf("CountWorkflowExecutions operation failed. Query failed: %v", err))
+		return nil, convertSQLError("CountWorkflowExecutions operation failed.", err)
 	}
 	resp := &manager.CountWorkflowExecutionsResponse{
 		Count:  0,
@@ -505,8 +507,7 @@ func (s *VisibilityStore) GetWorkflowExecution(
 		RunID:       request.RunID,
 	})
 	if err != nil {
-		return nil, serviceerror.NewUnavailable(
-			fmt.Sprintf("GetWorkflowExecution operation failed. Select failed: %v", err))
+		return nil, convertSQLError("GetWorkflowExecution operation failed.", err)
 	}
 	info, err := s.rowToInfo(row, request.Namespace)
 	if err != nil {

docs/architecture/retry.md

@@ -13,7 +13,7 @@ It's important to note that a special retry policy is used for `ResourceExhauste
 
 ## Service Error
 
-Service errors are specific Go errors that can generate a gRCP `Status` (see [status.proto](https://github.com/grpc/grpc/blob/master/src/proto/grpc/status/status.proto)). 
+Service errors are specific Go errors that can generate a gRPC `Status` (see [status.proto](https://github.com/grpc/grpc/blob/master/src/proto/grpc/status/status.proto)). 
 A gRPC status contains a gRPC `Code` (see [code.proto](https://github.com/googleapis/googleapis/blob/master/google/rpc/code.proto)), a message and (optionally) a payload with more details.
 
 ```go