chore: move timing aggregation and timing to outside the cmd exec

Emyrk · Emyrk · commit 6b8d46686749 · 2025-12-09T09:22:52.000-06:00
diff --git a/provisioner/terraform/executor.go b/provisioner/terraform/executor.go
@@ -325,9 +325,7 @@ func (e *executor) plan(ctx, killCtx context.Context, env, vars []string, logr l
 		<-doneErr
 	}()
 
-	endStage := e.timings.startStage(database.ProvisionerJobTimingStagePlan)
 	err := e.execWriteOutput(ctx, killCtx, args, env, outWriter, errWriter)
-	endStage(err)
 	if err != nil {
 		return nil, xerrors.Errorf("terraform plan: %w", err)
 	}
@@ -386,13 +384,11 @@ func (e *executor) plan(ctx, killCtx context.Context, env, vars []string, logr l
 		}
 	}
 
-	// DoneOut must be completed before we aggregate timings to ensure all logs have been processed.
-	<-doneOut
 	msg := &proto.PlanComplete{
 		Parameters:            state.Parameters,
 		Resources:             state.Resources,
 		ExternalAuthProviders: state.ExternalAuthProviders,
-		Timings:               append(e.timings.aggregate(), graphTimings.aggregate()...),
+		Timings:               graphTimings.aggregate(),
 		Presets:               state.Presets,
 		Plan:                  planJSON,
 		ResourceReplacements:  resReps,
@@ -601,9 +597,7 @@ func (e *executor) apply(
 	}()
 
 	// `terraform apply`
-	endStage := e.timings.startStage(database.ProvisionerJobTimingStageApply)
 	err := e.execWriteOutput(ctx, killCtx, args, env, outWriter, errWriter)
-	endStage(err)
 	if err != nil {
 		return nil, xerrors.Errorf("terraform apply: %w", err)
 	}
@@ -619,15 +613,11 @@ func (e *executor) apply(
 		return nil, xerrors.Errorf("read statefile %q: %w", statefilePath, err)
 	}
 
-	// DoneOut must be completed before we aggregate timings to ensure all logs have been processed.
-	<-doneOut
-	agg := e.timings.aggregate()
 	return &proto.ApplyComplete{
 		Parameters:            state.Parameters,
 		Resources:             state.Resources,
 		ExternalAuthProviders: state.ExternalAuthProviders,
 		State:                 stateContent,
-		Timings:               agg,
 		AiTasks:               state.AITasks,
 	}, nil
 }
diff --git a/provisioner/terraform/provision.go b/provisioner/terraform/provision.go
@@ -160,14 +160,16 @@ func (s *server) Plan(
 		return provisionersdk.PlanErrorf("plan vars: %s", err)
 	}
 
+	endPlanStage := e.timings.startStage(database.ProvisionerJobTimingStagePlan)
 	resp, err := e.plan(ctx, killCtx, env, vars, sess, request)
+	endPlanStage(err)
 	if err != nil {
 		return provisionersdk.PlanErrorf("%s", err.Error())
 	}
 
 	// Prepend init timings since they occur prior to plan timings.
 	// Order is irrelevant; this is merely indicative.
-	resp.Timings = append(initTimings.aggregate(), resp.Timings...) // mergeInitTimings(initTimings.aggregate(), resp.Timings)
+	resp.Timings = append(resp.Timings, append(initTimings.aggregate(), e.timings.aggregate()...)...)
 	resp.Modules = modules
 	return resp
 }
@@ -204,9 +206,11 @@ func (s *server) Apply(
 		return provisionersdk.ApplyErrorf("provision env: %s", err)
 	}
 	env = otelEnvInject(ctx, env)
+	endStage := e.timings.startStage(database.ProvisionerJobTimingStageApply)
 	resp, err := e.apply(
 		ctx, killCtx, env, sess,
 	)
+	endStage(err)
 	if err != nil {
 		errorMessage := err.Error()
 		// Terraform can fail and apply and still need to store it's state.
@@ -217,6 +221,7 @@ func (s *server) Apply(
 			Error: errorMessage,
 		}
 	}
+	resp.Timings = e.timings.aggregate()
 	return resp
 }
 

Original file line number	Diff line number	Diff line change
`@@ -160,14 +160,16 @@ func (s *server) Plan(`
`160`	`160`	`return provisionersdk.PlanErrorf("plan vars: %s", err)`
`161`	`161`	`}`
`162`	`162`
	`163`	`+ endPlanStage := e.timings.startStage(database.ProvisionerJobTimingStagePlan)`
`163`	`164`	`resp, err := e.plan(ctx, killCtx, env, vars, sess, request)`
	`165`	`+ endPlanStage(err)`
`164`	`166`	`if err != nil {`
`165`	`167`	`return provisionersdk.PlanErrorf("%s", err.Error())`
`166`	`168`	`}`
`167`	`169`
`168`	`170`	`// Prepend init timings since they occur prior to plan timings.`
`169`	`171`	`// Order is irrelevant; this is merely indicative.`
`170`		`- resp.Timings = append(initTimings.aggregate(), resp.Timings...) // mergeInitTimings(initTimings.aggregate(), resp.Timings)`
	`172`	`+ resp.Timings = append(resp.Timings, append(initTimings.aggregate(), e.timings.aggregate()...)...)`
`171`	`173`	`resp.Modules = modules`
`172`	`174`	`return resp`
`173`	`175`	`}`
`@@ -204,9 +206,11 @@ func (s *server) Apply(`
`204`	`206`	`return provisionersdk.ApplyErrorf("provision env: %s", err)`
`205`	`207`	`}`
`206`	`208`	`env = otelEnvInject(ctx, env)`
	`209`	`+ endStage := e.timings.startStage(database.ProvisionerJobTimingStageApply)`
`207`	`210`	`resp, err := e.apply(`
`208`	`211`	`ctx, killCtx, env, sess,`
`209`	`212`	`)`
	`213`	`+ endStage(err)`
`210`	`214`	`if err != nil {`
`211`	`215`	`errorMessage := err.Error()`
`212`	`216`	`// Terraform can fail and apply and still need to store it's state.`
`@@ -217,6 +221,7 @@ func (s *server) Apply(`
`217`	`221`	`Error: errorMessage,`
`218`	`222`	`}`
`219`	`223`	`}`
	`224`	`+ resp.Timings = e.timings.aggregate()`
`220`	`225`	`return resp`
`221`	`226`	`}`
`222`	`227`