Address PR review feedback (#874)

- bulk insert command queue messages within the transaction bound
- finish command batches after every enqueue failure
- build the outstanding queue index concurrently

Author: ankitgoswami

server/generated/sqlc/db.go

@@ -0,0 +1,79 @@
+package command_test
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	commonpb "github.com/block/proto-fleet/server/generated/grpc/common/v1"
+	pb "github.com/block/proto-fleet/server/generated/grpc/minercommand/v1"
+	"github.com/block/proto-fleet/server/internal/domain/commandtype"
+	"github.com/block/proto-fleet/server/internal/infrastructure/queue"
+	"github.com/block/proto-fleet/server/internal/testutil"
+)
+
+type failingEnqueueQueue struct {
+	err error
+}
+
+func (q failingEnqueueQueue) Enqueue(context.Context, string, commandtype.Type, []int64, interface{}) error {
+	return q.err
+}
+
+func (q failingEnqueueQueue) EnqueueMany(context.Context, string, commandtype.Type, []queue.EnqueueMessage) error {
+	return q.err
+}
+
+func (failingEnqueueQueue) Dequeue(context.Context, int32) ([]queue.Message, error) {
+	return nil, nil
+}
+
+func (failingEnqueueQueue) IsBatchFinished(context.Context, string) (bool, error) {
+	return false, nil
+}
+
+func (failingEnqueueQueue) MaxFailureRetries() int32 {
+	return 0
+}
+
+func TestCommandEnqueueFailureFinishesCreatedBatch(t *testing.T) {
+	if testing.Short() {
+		t.Skip("Skipping database integration test in short mode")
+	}
+
+	// Arrange
+	conn, dbService, user := setupRetentionTest(t)
+	device := dbService.CreateDevice(user.OrganizationID, "proto")
+	enqueueErr := errors.New("queue unavailable")
+	svc := newDispatchIntegrationTestService(t, conn, failingEnqueueQueue{err: enqueueErr})
+	ctx := testutil.MockAuthContextForTesting(t.Context(), user.DatabaseID, user.OrganizationID)
+
+	// Act
+	result, err := svc.BlinkLED(ctx, &pb.DeviceSelector{
+		SelectionType: &pb.DeviceSelector_IncludeDevices{
+			IncludeDevices: &commonpb.DeviceIdentifierList{DeviceIdentifiers: []string{device.ID}},
+		},
+	})
+
+	// Assert
+	require.Error(t, err)
+	assert.Nil(t, result)
+	assert.ErrorContains(t, err, enqueueErr.Error())
+	var batchUUID, status string
+	require.NoError(t, conn.QueryRowContext(t.Context(), `
+		SELECT uuid, status
+		FROM command_batch_log
+		WHERE organization_id = $1
+		ORDER BY id DESC
+		LIMIT 1
+	`, user.OrganizationID).Scan(&batchUUID, &status))
+	assert.Equal(t, "FINISHED", status)
+	var queued int
+	require.NoError(t, conn.QueryRowContext(t.Context(), `
+		SELECT COUNT(*) FROM queue_message WHERE command_batch_log_uuid = $1
+	`, batchUUID).Scan(&queued))
+	assert.Zero(t, queued)
+}

server/generated/sqlc/querier.go

@@ -1061,15 +1061,16 @@ func (s *Service) processCommand(ctx context.Context, command *Command) (*Comman
 		return s.messageQueue.EnqueueMany(workCtx, batchLogIdentifier, command.commandType, queuePayloads)
 	})
 	if err != nil {
+		var enqueueErr error
 		switch {
 		case errors.Is(err, errExecutionStoppedBeforeEnqueue):
-			enqueueErr := fleeterror.NewInternalError("command execution service stopped before enqueue")
-			return nil, s.finishUnenqueuedCommandBatch(ctx, batchLogIdentifier, enqueueErr)
+			enqueueErr = fleeterror.NewInternalError("command execution service stopped before enqueue")
 		case len(queuePayloads) == 0:
-			return nil, fleeterror.NewInternalErrorf("error enqueuing a batch of commands: %v", err)
+			enqueueErr = fleeterror.NewInternalErrorf("error enqueuing a batch of commands: %v", err)
 		default:
-			return nil, fleeterror.NewInternalErrorf("error enqueuing per-device command payloads: %v", err)
+			enqueueErr = fleeterror.NewInternalErrorf("error enqueuing per-device command payloads: %v", err)
 		}
+		return nil, s.finishUnenqueuedCommandBatch(ctx, batchLogIdentifier, enqueueErr)
 	}
 
 	return &CommandResult{
@@ -1548,14 +1549,10 @@ func (s *Service) ReapplyCurrentPoolsWithWorkerNames(
 		return s.enqueueWorkerNameReapplyMessages(workCtx, commandBatchLogUUID, deviceIdentifiers, deviceIDsByIdentifier, desiredWorkerNamesByDeviceIdentifier)
 	})
 	if err != nil {
-		enqueueErr := err
-		switch {
-		case errors.Is(err, errExecutionStoppedBeforeEnqueue):
-			enqueueErr = fleeterror.NewInternalError("command execution service stopped before enqueue")
-			return "", s.finishUnenqueuedCommandBatch(ctx, commandBatchLogUUID, enqueueErr)
-		default:
-			return "", enqueueErr
+		if errors.Is(err, errExecutionStoppedBeforeEnqueue) {
+			err = fleeterror.NewInternalError("command execution service stopped before enqueue")
 		}
+		return "", s.finishUnenqueuedCommandBatch(ctx, commandBatchLogUUID, err)
 	}
 
 	s.initializeStatusUpdateRoutine(commandBatchLogUUID, nil)

server/generated/sqlc/queue.sql.go

@@ -23,6 +23,14 @@ import (
 
 func newZeroTargetDispatchTestService(t *testing.T, conn *sql.DB) *command.Service {
 	t.Helper()
+	return newDispatchIntegrationTestService(t, conn, queue.NewDatabaseMessageQueue(&queue.Config{
+		DequeLimit:        10,
+		MaxFailureRetries: 1,
+	}, conn))
+}
+
+func newDispatchIntegrationTestService(t *testing.T, conn *sql.DB, messageQueue queue.MessageQueue) *command.Service {
+	t.Helper()
 
 	commandConfig := &command.Config{
 		MaxWorkers:                       1,
@@ -33,11 +41,6 @@ func newZeroTargetDispatchTestService(t *testing.T, conn *sql.DB) *command.Servi
 		StuckMessageTimeout:              time.Hour,
 		ReaperInterval:                   time.Hour,
 	}
-	queueConfig := &queue.Config{
-		DequeLimit:        10,
-		MaxFailureRetries: 1,
-	}
-	messageQueue := queue.NewDatabaseMessageQueue(queueConfig, conn)
 	executionCtx, cancel := context.WithCancel(context.Background())
 	t.Cleanup(cancel)
 

server/generated/sqlc/retrying_querier.gen.go

@@ -5,8 +5,6 @@ import (
 	"database/sql"
 	"encoding/json"
 
-	"github.com/sqlc-dev/pqtype"
-
 	"github.com/block/proto-fleet/server/generated/sqlc"
 	"github.com/block/proto-fleet/server/internal/domain/commandtype"
 	"github.com/block/proto-fleet/server/internal/domain/fleeterror"
@@ -58,19 +56,20 @@ func (d DatabaseMessageQueue) EnqueueMany(ctx context.Context, commandBatchLogUU
 }
 
 func (d DatabaseMessageQueue) enqueueEncoded(ctx context.Context, commandBatchLogUUID string, commandType commandtype.Type, messages []encodedMessage) error {
+	deviceIDs := make([]int64, len(messages))
+	payloads := make([]string, len(messages))
+	for i, message := range messages {
+		deviceIDs[i] = message.deviceID
+		payloads[i] = string(message.payload)
+	}
 	return db.WithTransactionTimeoutNoResult(ctx, d.conn, runtimepolicy.CommandTransactionBound, func(q sqlc.Querier) error {
-		for _, message := range messages {
-			err := q.CreateQueueMessage(ctx, sqlc.CreateQueueMessageParams{
-				CommandBatchLogUuid: commandBatchLogUUID,
-				CommandType:         commandType.String(),
-				DeviceID:            message.deviceID,
-				Status:              sqlc.QueueStatusEnumPENDING,
-				RetryCount:          0,
-				Payload:             pqtype.NullRawMessage{RawMessage: message.payload, Valid: true},
-			})
-			if err != nil {
-				return fleeterror.NewInternalErrorf("failed to enqueue message: %v", err)
-			}
+		if err := q.CreateQueueMessages(ctx, sqlc.CreateQueueMessagesParams{
+			CommandBatchLogUuid: commandBatchLogUUID,
+			CommandType:         commandType.String(),
+			DeviceIds:           deviceIDs,
+			Payloads:            payloads,
+		}); err != nil {
+			return fleeterror.NewInternalErrorf("failed to enqueue messages: %v", err)
 		}
 		return nil
 	})

server/internal/domain/command/enqueue_failure_integration_test.go

@@ -0,0 +1,79 @@
+package queue_test
+
+import (
+	"database/sql"
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/sqlc-dev/pqtype"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/block/proto-fleet/server/generated/sqlc"
+	"github.com/block/proto-fleet/server/internal/domain/commandtype"
+	"github.com/block/proto-fleet/server/internal/infrastructure/id"
+	"github.com/block/proto-fleet/server/internal/infrastructure/queue"
+	"github.com/block/proto-fleet/server/internal/testutil"
+)
+
+func TestDatabaseMessageQueueEnqueueManyInsertsPerDevicePayloads(t *testing.T) {
+	if testing.Short() {
+		t.Skip("Skipping database integration test in short mode")
+	}
+
+	// Arrange
+	cfg, err := testutil.GetTestConfig()
+	require.NoError(t, err)
+	dbService := testutil.NewDatabaseService(t, cfg)
+	user := dbService.CreateSuperAdminUser()
+	firstDevice := dbService.CreateDevice(user.OrganizationID, "proto")
+	secondDevice := dbService.CreateDevice(user.OrganizationID, "proto")
+	batchUUID := id.GenerateID()
+	commandType := commandtype.UpdateMiningPools
+	_, err = sqlc.New(dbService.DB).CreateCommandBatchLog(t.Context(), sqlc.CreateCommandBatchLogParams{
+		Uuid:           batchUUID,
+		Type:           commandType.String(),
+		CreatedBy:      user.DatabaseID,
+		CreatedAt:      time.Now(),
+		Status:         sqlc.BatchStatusEnumPENDING,
+		DevicesCount:   2,
+		Payload:        pqtype.NullRawMessage{},
+		OrganizationID: sql.NullInt64{Int64: user.OrganizationID, Valid: true},
+	})
+	require.NoError(t, err)
+	messageQueue := queue.NewDatabaseMessageQueue(&queue.Config{}, dbService.DB)
+	messages := []queue.EnqueueMessage{
+		{DeviceID: firstDevice.DatabaseID, Payload: map[string]string{"worker_name": "first"}},
+		{DeviceID: secondDevice.DatabaseID, Payload: map[string]string{"worker_name": "second"}},
+	}
+
+	// Act
+	err = messageQueue.EnqueueMany(t.Context(), batchUUID, commandType, messages)
+
+	// Assert
+	require.NoError(t, err)
+	rows, err := dbService.DB.QueryContext(t.Context(), `
+		SELECT device_id, payload, status
+		FROM queue_message
+		WHERE command_batch_log_uuid = $1
+	`, batchUUID)
+	require.NoError(t, err)
+	defer rows.Close()
+	gotPayloads := make(map[int64]map[string]string)
+	for rows.Next() {
+		var deviceID int64
+		var payload []byte
+		var status string
+		require.NoError(t, rows.Scan(&deviceID, &payload, &status))
+		var decoded map[string]string
+		require.NoError(t, json.Unmarshal(payload, &decoded))
+		gotPayloads[deviceID] = decoded
+		assert.Equal(t, "PENDING", status)
+	}
+	require.NoError(t, rows.Err())
+	assert.Equal(t, map[int64]map[string]string{
+		firstDevice.DatabaseID:  {"worker_name": "first"},
+		secondDevice.DatabaseID: {"worker_name": "second"},
+	}, gotPayloads)
+}

server/internal/domain/command/service.go

@@ -1 +1 @@
-DROP INDEX IF EXISTS idx_queue_message_reaper_outstanding;
+DROP INDEX CONCURRENTLY IF EXISTS idx_queue_message_reaper_outstanding;