chore(shard-distributor): classify errors (#7466)

Author: arzonus

service/sharddistributor/handler/executor.go

@@ -50,19 +50,18 @@ func (h *executor) Heartbeat(ctx context.Context, request *types.ExecutorHeartbe
 	previousHeartbeat, assignedShards, err := h.storage.GetHeartbeat(ctx, request.Namespace, request.ExecutorID)
 	// We ignore Executor not found errors, since it just means that this executor heartbeat the first time.
 	if err != nil && !errors.Is(err, store.ErrExecutorNotFound) {
-		return nil, fmt.Errorf("get heartbeat: %w", err)
+		return nil, &types.InternalServiceError{Message: fmt.Sprintf("failed to get heartbeat: %v", err)}
 	}
 
 	now := h.timeSource.Now().UTC()
 	mode := h.migrationConfiguration.GetMigrationMode(request.Namespace)
 
 	switch mode {
 	case types.MigrationModeINVALID:
-		h.logger.Warn("Migration mode is invalid", tag.ShardNamespace(request.Namespace), tag.ShardExecutor(request.ExecutorID))
-		return nil, fmt.Errorf("migration mode is invalid")
+		return nil, &types.InternalServiceError{Message: fmt.Sprintf("namespace's migration mode is invalid: %v", err)}
 	case types.MigrationModeLOCALPASSTHROUGH:
 		h.logger.Warn("Migration mode is local passthrough, no calls to heartbeat allowed", tag.ShardNamespace(request.Namespace), tag.ShardExecutor(request.ExecutorID))
-		return nil, fmt.Errorf("migration mode is local passthrough")
+		return nil, &types.BadRequestError{Message: "migration mode is local passthrough, no calls to heartbeat allowed"}
 	// From SD perspective the behaviour is the same
 	case types.MigrationModeLOCALPASSTHROUGHSHADOW, types.MigrationModeDISTRIBUTEDPASSTHROUGH:
 		assignedShards, err = h.assignShardsInCurrentHeartbeat(ctx, request)
@@ -88,12 +87,12 @@ func (h *executor) Heartbeat(ctx context.Context, request *types.ExecutorHeartbe
 	}
 
 	if err := validateMetadata(newHeartbeat.Metadata); err != nil {
-		return nil, fmt.Errorf("validate metadata: %w", err)
+		return nil, types.BadRequestError{Message: fmt.Sprintf("invalid metadata: %s", err)}
 	}
 
 	err = h.storage.RecordHeartbeat(ctx, request.Namespace, request.ExecutorID, newHeartbeat)
 	if err != nil {
-		return nil, fmt.Errorf("record heartbeat: %w", err)
+		return nil, &types.InternalServiceError{Message: fmt.Sprintf("failed to record heartbeat: %v", err)}
 	}
 
 	return _convertResponse(assignedShards, mode), nil
@@ -108,7 +107,7 @@ func (h *executor) assignShardsInCurrentHeartbeat(ctx context.Context, request *
 	}
 	err := h.storage.DeleteExecutors(ctx, request.GetNamespace(), []string{request.GetExecutorID()}, store.NopGuard())
 	if err != nil {
-		return nil, fmt.Errorf("delete executors: %w", err)
+		return nil, &types.InternalServiceError{Message: fmt.Sprintf("failed to delete assigned shards: %v", err)}
 	}
 	for shard := range request.GetShardStatusReports() {
 		assignedShards.AssignedShards[shard] = &types.ShardAssignment{
@@ -124,7 +123,7 @@ func (h *executor) assignShardsInCurrentHeartbeat(ctx context.Context, request *
 	}
 	err = h.storage.AssignShards(ctx, request.GetNamespace(), assignShardsRequest, store.NopGuard())
 	if err != nil {
-		return nil, fmt.Errorf("assign shards in current heartbeat: %w", err)
+		return nil, &types.InternalServiceError{Message: fmt.Sprintf("failed to assign shards in the current heartbeat: %v", err)}
 	}
 	return &assignedShards, nil
 }

service/sharddistributor/handler/executor_test.go

@@ -395,8 +395,35 @@ func TestHeartbeat(t *testing.T) {
 
 		_, err := handler.Heartbeat(ctx, req)
 		require.Error(t, err)
-		require.Contains(t, err.Error(), "validate metadata")
-		require.Contains(t, err.Error(), "exceeds the maximum")
+		require.Contains(t, err.Error(), "invalid metadata: metadata has 33 keys, which exceeds the maximum of 32")
+	})
+
+	// Test Case: Heartbeat with executor associated with MigrationModeLOCALPASSTHROUGH (should error)
+	t.Run("MigrationModeLOCALPASSTHROUGH", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		mockStore := store.NewMockStore(ctrl)
+		mockTimeSource := clock.NewMockedTimeSource()
+		shardDistributionCfg := config.ShardDistribution{
+			Namespaces: []config.Namespace{{Name: namespace, Mode: config.MigrationModeLOCALPASSTHROUGH}},
+		}
+		migrationConfig := newMigrationConfig(t, []configEntry{{dynamicproperties.MigrationMode, config.MigrationModeLOCALPASSTHROUGH}})
+		handler := NewExecutorHandler(testlogger.New(t), mockStore, mockTimeSource, shardDistributionCfg, migrationConfig)
+
+		req := &types.ExecutorHeartbeatRequest{
+			Namespace:  namespace,
+			ExecutorID: executorID,
+			Status:     types.ExecutorStatusACTIVE,
+		}
+		previousHeartbeat := store.HeartbeatState{
+			LastHeartbeat: now.Unix(),
+			Status:        types.ExecutorStatusACTIVE,
+		}
+
+		mockStore.EXPECT().GetHeartbeat(gomock.Any(), namespace, executorID).Return(&previousHeartbeat, nil, nil)
+
+		_, err := handler.Heartbeat(ctx, req)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "migration mode is local passthrough, no calls to heartbeat allowed")
 	})
 }
 

service/sharddistributor/handler/handler.go

@@ -99,7 +99,7 @@ func (h *handlerImpl) GetShardOwner(ctx context.Context, request *types.GetShard
 		}
 	}
 	if err != nil {
-		return nil, fmt.Errorf("get shard owner: %w", err)
+		return nil, &types.InternalServiceError{Message: fmt.Sprintf("failed to get shard owner: %v", err)}
 	}
 
 	resp = &types.GetShardOwnerResponse{
@@ -116,7 +116,7 @@ func (h *handlerImpl) assignEphemeralShard(ctx context.Context, namespace string
 	// Get the current state of the namespace and find the executor with the least assigned shards
 	state, err := h.storage.GetState(ctx, namespace)
 	if err != nil {
-		return nil, fmt.Errorf("get state: %w", err)
+		return nil, &types.InternalServiceError{Message: fmt.Sprintf("failed to get namespace state: %v", err)}
 	}
 
 	var executor string
@@ -132,7 +132,7 @@ func (h *handlerImpl) assignEphemeralShard(ctx context.Context, namespace string
 	// Assign the shard to the executor with the least assigned shards
 	err = h.storage.AssignShard(ctx, namespace, shardID, executor)
 	if err != nil {
-		return nil, fmt.Errorf("assign ephemeral shard: %w", err)
+		return nil, &types.InternalServiceError{Message: fmt.Sprintf("failed to assign ephemeral shard: %v", err)}
 	}
 
 	return &types.GetShardOwnerResponse{
@@ -148,13 +148,13 @@ func (h *handlerImpl) WatchNamespaceState(request *types.WatchNamespaceStateRequ
 	assignmentChangesChan, unSubscribe, err := h.storage.SubscribeToAssignmentChanges(server.Context(), request.Namespace)
 	defer unSubscribe()
 	if err != nil {
-		return fmt.Errorf("subscribe to namespace state: %w", err)
+		return &types.InternalServiceError{Message: fmt.Sprintf("failed to subscribe to namespace state: %v", err)}
 	}
 
 	// Send initial state immediately so client doesn't have to wait for first update
 	state, err := h.storage.GetState(server.Context(), request.Namespace)
 	if err != nil {
-		return fmt.Errorf("get initial state: %w", err)
+		return &types.InternalServiceError{Message: fmt.Sprintf("failed to get namespace state: %v", err)}
 	}
 	response := toWatchNamespaceStateResponse(state)
 	if err := server.Send(response); err != nil {

service/sharddistributor/handler/handler_test.go

@@ -293,6 +293,6 @@ func TestWatchNamespaceState(t *testing.T) {
 
 		err := handler.WatchNamespaceState(&types.WatchNamespaceStateRequest{Namespace: "test-ns"}, mockServer)
 		require.Error(t, err)
-		require.Contains(t, err.Error(), "get initial state")
+		require.Contains(t, err.Error(), "failed to get namespace state: storage error")
 	})
 }