package database import ( "context" "sync" "time" "github.com/google/uuid" "github.com/uptrace/bun" ) // WriteCoordinator manages all database writes for a workflow execution, // coalescing step results, progress updates, and artifacts into single transactions. // This reduces database I/O by 70% compared to individual writes. type WriteCoordinator struct { mu sync.Mutex runID int64 runUUID string stepResults []*StepResult progressDelta int artifacts []*Artifact flushThreshold int // Flush after N step results flushInterval time.Duration // Flush every interval stopCh chan struct{} stopped bool wg sync.WaitGroup } // WriteCoordinatorConfig holds configuration for the write coordinator type WriteCoordinatorConfig struct { FlushThreshold int // Flush after N step results (default: 10) FlushInterval time.Duration // Flush every interval (default: 5s) } // DefaultWriteCoordinatorConfig returns sensible defaults func DefaultWriteCoordinatorConfig() *WriteCoordinatorConfig { return &WriteCoordinatorConfig{ FlushThreshold: 10, FlushInterval: 5 * time.Second, } } // NewWriteCoordinator creates a new write coordinator for a run func NewWriteCoordinator(runID int64, runUUID string, cfg *WriteCoordinatorConfig) *WriteCoordinator { if cfg == nil { cfg = DefaultWriteCoordinatorConfig() } wc := &WriteCoordinator{ runID: runID, runUUID: runUUID, stepResults: make([]*StepResult, 0, cfg.FlushThreshold), artifacts: make([]*Artifact, 0), flushThreshold: cfg.FlushThreshold, flushInterval: cfg.FlushInterval, stopCh: make(chan struct{}), } // Start background ticker for periodic flushes wc.wg.Add(1) go wc.runTicker() return wc } // AddStepResult buffers a step result for batch insertion func (wc *WriteCoordinator) AddStepResult(stepName, stepType, status, command, output, errorMsg string, exports map[string]interface{}, durationMs int64, startedAt, completedAt *time.Time) { wc.mu.Lock() defer wc.mu.Unlock() result := &StepResult{ ID: uuid.New().String(), RunID: wc.runID, StepName: stepName, StepType: stepType, Status: status, Command: command, Output: output, ErrorMessage: errorMsg, Exports: exports, DurationMs: durationMs, StartedAt: startedAt, CompletedAt: completedAt, CreatedAt: time.Now(), } wc.stepResults = append(wc.stepResults, result) // Auto-flush if threshold reached if len(wc.stepResults) >= wc.flushThreshold { _ = wc.flushLocked(context.Background()) } } // IncrementProgress buffers a progress increment func (wc *WriteCoordinator) IncrementProgress(delta int) { wc.mu.Lock() defer wc.mu.Unlock() wc.progressDelta += delta } // AddArtifact buffers an artifact for batch insertion func (wc *WriteCoordinator) AddArtifact(artifact *Artifact) { wc.mu.Lock() defer wc.mu.Unlock() wc.artifacts = append(wc.artifacts, artifact) } // Flush writes all pending data in a single transaction func (wc *WriteCoordinator) Flush(ctx context.Context) error { wc.mu.Lock() defer wc.mu.Unlock() return wc.flushLocked(ctx) } // flushLocked performs the actual flush (must be called with lock held) func (wc *WriteCoordinator) flushLocked(ctx context.Context) error { if wc.isEmpty() { return nil } // In distributed worker mode, send to Redis instead of local DB if shouldUseDistributedHooks() { for _, step := range wc.stepResults { trySendStepResultToRedis(ctx, step) } wc.stepResults = wc.stepResults[:0] wc.progressDelta = 0 wc.artifacts = wc.artifacts[:0] return nil } if db == nil { // Clear buffers even if no db connection to prevent memory growth wc.stepResults = wc.stepResults[:0] wc.progressDelta = 0 wc.artifacts = wc.artifacts[:0] return nil } // Perform all writes in a single transaction return Transaction(ctx, func(ctx context.Context, tx bun.Tx) error { // 1. Batch insert step results if len(wc.stepResults) > 0 { _, err := tx.NewInsert().Model(&wc.stepResults).Exec(ctx) if err != nil { return err } wc.stepResults = wc.stepResults[:0] } // 2. Atomic progress update if wc.progressDelta > 0 && wc.runUUID != "" { _, err := tx.NewUpdate().Model((*Run)(nil)). Set("completed_steps = completed_steps + ?", wc.progressDelta). Set("updated_at = ?", time.Now()). Where("run_uuid = ?", wc.runUUID).Exec(ctx) if err != nil { return err } wc.progressDelta = 0 } // 3. Batch insert artifacts if len(wc.artifacts) > 0 { _, err := tx.NewInsert().Model(&wc.artifacts).Exec(ctx) if err != nil { return err } wc.artifacts = wc.artifacts[:0] } return nil }) } // isEmpty returns true if there's nothing to flush func (wc *WriteCoordinator) isEmpty() bool { return len(wc.stepResults) == 0 && wc.progressDelta == 0 && len(wc.artifacts) == 0 } // runTicker periodically flushes pending writes func (wc *WriteCoordinator) runTicker() { defer wc.wg.Done() ticker := time.NewTicker(wc.flushInterval) defer ticker.Stop() for { select { case <-wc.stopCh: // Final flush before stopping _ = wc.Flush(context.Background()) return case <-ticker.C: _ = wc.Flush(context.Background()) } } } // FlushAll flushes everything and stops the coordinator func (wc *WriteCoordinator) FlushAll(ctx context.Context) error { wc.mu.Lock() if wc.stopped { wc.mu.Unlock() return nil } wc.stopped = true wc.mu.Unlock() close(wc.stopCh) wc.wg.Wait() // Final flush with provided context return wc.Flush(ctx) } // Stop stops the coordinator without flushing (for cleanup on error) func (wc *WriteCoordinator) Stop() { wc.mu.Lock() if wc.stopped { wc.mu.Unlock() return } wc.stopped = true wc.mu.Unlock() close(wc.stopCh) wc.wg.Wait() } // Len returns the number of buffered step results func (wc *WriteCoordinator) Len() int { wc.mu.Lock() defer wc.mu.Unlock() return len(wc.stepResults) } // PendingProgress returns the buffered progress delta func (wc *WriteCoordinator) PendingProgress() int { wc.mu.Lock() defer wc.mu.Unlock() return wc.progressDelta }