Update execution server to allow storing metadata in Redis (#9357)

Introduces 3 flags to the execution server:
- `remote_execution.write_execution_progress_state_to_redis`: if
enabled, writes execution updates and invocation-execution links to
Redis, using the new methods added in
9b6419d
- `remote_execution.write_executions_to_primary_db`: if enabled, writes
Execution rows and InvocationExecution rows to mysql. Defaults to true
in order to preserve our current behavior.
- `remote_execution.read_final_execution_state_from_redis`: if enabled,
read the final execution state from redis instead of mysql when we write
executions to clickhouse.

The rollout plan to fully stop writing executions to MySQL would then
be:
- Rollout 1: enable
`remote_execution.write_execution_progress_state_to_redis`
- Rollout 2: enable
`remote_execution.read_final_execution_state_from_redis`
- Rollout 3: disable `remote_execution.write_executions_to_primary_db`

This PR only affects the execution server; the UI endpoint
(execution_service) also needs to be updated to read the in-progress
execution state from redis, which will be done in a separate PR.

Author: bduffany

enterprise/server/remote_execution/execution_server/execution_server.go

@@ -16,7 +16,6 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/operation"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/platform"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/tasksize"
-	"github.com/buildbuddy-io/buildbuddy/enterprise/server/util/execution"
 	"github.com/buildbuddy-io/buildbuddy/proto/invocation_status"
 	"github.com/buildbuddy-io/buildbuddy/server/environment"
 	"github.com/buildbuddy-io/buildbuddy/server/interfaces"
@@ -48,6 +47,7 @@ import (
 	"google.golang.org/protobuf/types/known/anypb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
+	executil "github.com/buildbuddy-io/buildbuddy/enterprise/server/util/execution"
 	espb "github.com/buildbuddy-io/buildbuddy/proto/execution_stats"
 	repb "github.com/buildbuddy-io/buildbuddy/proto/remote_execution"
 	scpb "github.com/buildbuddy-io/buildbuddy/proto/scheduler"
@@ -78,6 +78,10 @@ const (
 var (
 	enableRedisAvailabilityMonitoring = flag.Bool("remote_execution.enable_redis_availability_monitoring", false, "If enabled, the execution server will detect if Redis has lost state and will ask Bazel to retry executions.")
 	sharedExecutorPoolTeeRate         = flag.Float64("remote_execution.shared_executor_pool_tee_rate", 0, "If non-zero, work for the default shared executor pool will be teed to a separate experiment pool at this rate.", flag.Internal)
+
+	writeExecutionProgressStateToRedis = flag.Bool("remote_execution.write_execution_progress_state_to_redis", false, "If enabled, write initial execution metadata and progress updates (stage changes) to redis. This state is cleared when the execution is complete.", flag.Internal)
+	writeExecutionsToPrimaryDB         = flag.Bool("remote_execution.write_executions_to_primary_db", true, "If enabled, write executions and invocation-execution links to the primary DB.", flag.Internal)
+	readFinalExecutionStateFromRedis   = flag.Bool("remote_execution.read_final_execution_state_from_redis", false, "If enabled, read execution metadata and progress state from redis when writing executions to ClickHouse, instead of from the primary DB.", flag.Internal)
 )
 
 func fillExecutionFromActionMetadata(md *repb.ExecutedActionMetadata, execution *tables.Execution) {
@@ -219,7 +223,20 @@ func (s *ExecutionServer) insertExecution(ctx context.Context, executionID, invo
 	execution.GroupID = permissions.GroupID
 	execution.Perms = execution.Perms | permissions.Perms
 
-	return s.env.GetDBHandle().NewQuery(ctx, "execution_server_create_execution").Create(execution)
+	if *writeExecutionProgressStateToRedis {
+		now := time.Now()
+		execution.Model.CreatedAtUsec = now.UnixMicro()
+		execution.Model.UpdatedAtUsec = now.UnixMicro()
+		if err := s.env.GetExecutionCollector().UpdateInProgressExecution(ctx, executil.TableExecToProto(execution, nil /*=invocationLink*/)); err != nil {
+			log.CtxErrorf(ctx, "Failed to write execution update to redis: %s", err)
+		}
+	}
+
+	if *writeExecutionsToPrimaryDB {
+		return s.env.GetDBHandle().NewQuery(ctx, "execution_server_create_execution").Create(execution)
+	}
+
+	return nil
 }
 
 func (s *ExecutionServer) insertInvocationLink(ctx context.Context, executionID, invocationID string, linkType sipb.StoredInvocationLink_Type) error {
@@ -241,6 +258,10 @@ func (s *ExecutionServer) insertInvocationLink(ctx context.Context, executionID,
 		log.CtxWarningf(ctx, "failed to add invocation link(exeuction_id: %q invocation_id: %q, link_type: %d) in redis", executionID, invocationID, linkType)
 	}
 
+	if !*writeExecutionsToPrimaryDB {
+		return nil
+	}
+
 	link := &tables.InvocationExecution{
 		InvocationID: invocationID,
 		ExecutionID:  executionID,
@@ -265,7 +286,11 @@ func (s *ExecutionServer) insertInvocationLinkInRedis(ctx context.Context, execu
 		ExecutionId:  executionID,
 		Type:         linkType,
 	}
-	return s.env.GetExecutionCollector().AddExecutionInvocationLink(ctx, link, false /*=storeReverseLink*/)
+	// Only store the invocation => execution link if we're also writing
+	// execution progress state to redis, since these links are only used to
+	// list the in-progress state by invocation ID.
+	storeInvocationExecutionLink := *writeExecutionProgressStateToRedis
+	return s.env.GetExecutionCollector().AddExecutionInvocationLink(ctx, link, storeInvocationExecutionLink)
 }
 
 func trimStatus(statusMessage string) string {
@@ -315,6 +340,17 @@ func (s *ExecutionServer) updateExecution(ctx context.Context, executionID strin
 		}
 	}
 
+	if *writeExecutionProgressStateToRedis {
+		execution.Model.UpdatedAtUsec = time.Now().UnixMicro()
+		executionProto := executil.TableExecToProto(execution, nil /*=invocationLink*/)
+		if err := s.env.GetExecutionCollector().UpdateInProgressExecution(ctx, executionProto); err != nil {
+			log.CtxErrorf(ctx, "Failed to write execution update to redis: %s", err)
+		}
+	}
+
+	if !*writeExecutionsToPrimaryDB {
+		return nil
+	}
 	result := s.env.GetDBHandle().GORM(ctx, "execution_server_update_execution").Where(
 		"execution_id = ? AND stage != ?", executionID, repb.ExecutionStage_COMPLETED).Updates(execution)
 	dbErr := result.Error
@@ -349,17 +385,38 @@ func (s *ExecutionServer) recordExecution(
 	if s.env.GetDBHandle() == nil {
 		return status.FailedPreconditionError("database not configured")
 	}
-	var executionPrimaryDB tables.Execution
-
-	if err := s.env.GetDBHandle().NewQuery(ctx, "execution_server_lookup_execution").Raw(
-		`SELECT * FROM "Executions" WHERE execution_id = ?`, executionID).Take(&executionPrimaryDB); err != nil {
-		return status.InternalErrorf("failed to look up execution %q: %s", executionID, err)
+	var executionProto *repb.StoredExecution
+	// Read final execution details (including the initial metadata and
+	// completed execution result) from either redis or the primary DB. This
+	// state does not have the invocation details populated.
+	if *readFinalExecutionStateFromRedis {
+		ex, err := s.env.GetExecutionCollector().GetInProgressExecution(ctx, executionID)
+		if err != nil {
+			return status.InternalErrorf("failed to get execution %q from redis: %s", executionID, err)
+		}
+		if ex == nil {
+			return status.NotFoundErrorf("execution %q not found in redis", executionID)
+		}
+		executionProto = ex
+	} else {
+		var executionPrimaryDB tables.Execution
+		if err := s.env.GetDBHandle().NewQuery(ctx, "execution_server_lookup_execution").Raw(
+			`SELECT * FROM "Executions" WHERE execution_id = ?`, executionID).Take(&executionPrimaryDB); err != nil {
+			return status.InternalErrorf("failed to look up execution %q: %s", executionID, err)
+		}
+		executionProto = executil.TableExecToProto(&executionPrimaryDB, nil)
 	}
-	// Always clean up invocationLinks in Collector because we are not retrying
+	// Always clean up invocationLinks and execution updates from the collector.
+	// The execution cannot be retried after this point, so nothing will clean
+	// up this data if we don't do it here.
 	defer func() {
 		err := s.env.GetExecutionCollector().DeleteExecutionInvocationLinks(ctx, executionID)
 		if err != nil {
-			log.CtxErrorf(ctx, "failed to clean up invocation links in collector: %s", err)
+			log.CtxErrorf(ctx, "Failed to clean up invocation links in collector: %s", err)
+		}
+		err = s.env.GetExecutionCollector().DeleteInProgressExecution(ctx, executionID)
+		if err != nil {
+			log.CtxErrorf(ctx, "Failed to clean up in-progress execution in collector: %s", err)
 		}
 	}()
 	links, err := s.env.GetExecutionCollector().GetExecutionInvocationLinks(ctx, executionID)
@@ -368,7 +425,9 @@ func (s *ExecutionServer) recordExecution(
 	}
 	rmd := bazel_request.GetRequestMetadata(ctx)
 	for _, link := range links {
-		executionProto := execution.TableExecToProto(&executionPrimaryDB, link)
+		executionProto := executionProto.CloneVT()
+		executil.SetInvocationLink(executionProto, link)
+
 		// Set fields that aren't stored in the primary DB
 		executionProto.TargetLabel = rmd.GetTargetId()
 		executionProto.ActionMnemonic = rmd.GetActionMnemonic()

enterprise/server/remote_execution/execution_server/execution_server_test.go

@@ -399,6 +399,25 @@ func TestExecuteAndPublishOperation(t *testing.T) {
 			expectedExecutionUsage: tables.UsageCounts{LinuxExecutionDurationUsec: durationUsec},
 			publishMoreMetadata:    true,
 		},
+		{
+			name:                   "PublishMoreMetadata_PrimaryDBAndRedisDoubleWrite",
+			expectedExecutionUsage: tables.UsageCounts{LinuxExecutionDurationUsec: durationUsec},
+			publishMoreMetadata:    true,
+			flagOverrides: map[string]any{
+				"remote_execution.write_execution_progress_state_to_redis": true,
+				"remote_execution.write_executions_to_primary_db":          true,
+			},
+		},
+		{
+			name:                   "PublishMoreMetadata_NoPrimaryDB_RedisOnly",
+			expectedExecutionUsage: tables.UsageCounts{LinuxExecutionDurationUsec: durationUsec},
+			publishMoreMetadata:    true,
+			flagOverrides: map[string]any{
+				"remote_execution.write_execution_progress_state_to_redis": true,
+				"remote_execution.write_executions_to_primary_db":          false,
+				"remote_execution.read_final_execution_state_from_redis":   true,
+			},
+		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
 			testExecuteAndPublishOperation(t, test)
@@ -409,6 +428,7 @@ func TestExecuteAndPublishOperation(t *testing.T) {
 type publishTest struct {
 	name                     string
 	platformOverrides        map[string]string
+	flagOverrides            map[string]any
 	expectedSelfHosted       bool
 	expectedExecutionUsage   tables.UsageCounts
 	cachedResult, doNotCache bool
@@ -420,6 +440,9 @@ type publishTest struct {
 func testExecuteAndPublishOperation(t *testing.T, test publishTest) {
 	ctx := context.Background()
 	flags.Set(t, "app.enable_write_executions_to_olap_db", true)
+	for k, v := range test.flagOverrides {
+		flags.Set(t, k, v)
+	}
 	env, conn := setupEnv(t)
 	redis := testredis.Start(t)
 	env.SetDefaultRedisClient(redis.Client())

enterprise/server/util/execution/execution.go

@@ -19,12 +19,10 @@ import (
 )
 
 func TableExecToProto(in *tables.Execution, invLink *sipb.StoredInvocationLink) *repb.StoredExecution {
-	return &repb.StoredExecution{
+	ex := &repb.StoredExecution{
 		GroupId:                            in.GroupID,
 		UpdatedAtUsec:                      in.UpdatedAtUsec,
 		ExecutionId:                        in.ExecutionID,
-		InvocationUuid:                     strings.Replace(invLink.GetInvocationId(), "-", "", -1),
-		InvocationLinkType:                 int32(invLink.GetType()),
 		CreatedAtUsec:                      in.CreatedAtUsec,
 		UserId:                             in.UserID,
 		Worker:                             in.Worker,
@@ -55,6 +53,13 @@ func TableExecToProto(in *tables.Execution, invLink *sipb.StoredInvocationLink)
 		DoNotCache:                         in.DoNotCache,
 		CommandSnippet:                     in.CommandSnippet,
 	}
+	SetInvocationLink(ex, invLink)
+	return ex
+}
+
+func SetInvocationLink(ex *repb.StoredExecution, invLink *sipb.StoredInvocationLink) {
+	ex.InvocationUuid = strings.Replace(invLink.GetInvocationId(), "-", "", -1)
+	ex.InvocationLinkType = int32(invLink.GetType())
 }
 
 func TableExecToClientProto(in *tables.Execution) (*espb.Execution, error) {