【架构优化】从 SaveWorkspaceCheckpointStream 中拆出分片暂存写入 stageWorkspaceCheckpointChunks

internal/repo/workspace_checkpoint_chunks.go

@@ -56,60 +56,71 @@ func (r *RunsRepo) SaveWorkspaceCheckpointStream(runID, token string, now time.T
 		r.DB.Model(&WorkspaceCheckpoint{}).Select("COALESCE(archive_id, '')").Where("run_id = ?", runID)).Delete(&workspaceCheckpointChunk{}).Error; err != nil {
 		return err
 	}
-	hash := sha256.New()
+	size, count, sum, err := r.stageWorkspaceCheckpointChunks(runID, token, id, source)
+	if err != nil {
+		return err
+	}
+	if size == 0 {
+		return fmt.Errorf("empty workspace checkpoint")
+	}
+	err = r.DB.Transaction(func(tx *gorm.DB) error {
+		if _, err := r.takeRunForUpdate(tx, runID); err != nil {
+			return err
+		}
+		result := tx.Model(&WorkspaceCheckpoint{}).Where("run_id = ? AND lease_token = ? AND lease_until > ?", runID, token, time.Now().UTC()).Updates(map[string]any{
+			"archive": nil, "archive_id": id, "archive_bytes": size, "archive_chunks": count,
+			"sha256": sum, "saved_at": now, "error": "", "save_failures": 0, "save_retry_at": nil,
+		})
+		if result.Error != nil {
+			return result.Error
+		}
+		if result.RowsAffected != 1 {
+			return ErrWorkspaceInUse
+		}
+		return tx.Where("run_id = ? AND archive_id <> ?", runID, id).Delete(&workspaceCheckpointChunk{}).Error
+	})
+	published = err == nil
+	return err
+}
+
+// stageWorkspaceCheckpointChunks copies source into bounded chunk rows of
+// archiveID, re-checking the lease before every row. It returns the staged
+// byte size, chunk count and hex SHA-256; publishing is left to the caller.
+func (r *RunsRepo) stageWorkspaceCheckpointChunks(runID, token, archiveID string, source io.Reader) (int64, int, string, error) {
+	digest := sha256.New()
 	var size int64
 	count := 0
 	buffer := make([]byte, workspaceCheckpointChunkBytes)
 	for {
 		n, err := io.ReadFull(source, buffer)
 		if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
-			return err
+			return 0, 0, "", err
 		}
 		if n == 0 {
 			break
 		}
 		size += int64(n)
 		if size > workspaceCheckpointArchiveBytes {
-			return fmt.Errorf("workspace checkpoint exceeds archive size limit")
+			return 0, 0, "", fmt.Errorf("workspace checkpoint exceeds archive size limit")
 		}
 		var active int64
 		if err := r.DB.Model(&WorkspaceCheckpoint{}).Where("run_id = ? AND lease_token = ? AND lease_until > ?", runID, token, time.Now().UTC()).Count(&active).Error; err != nil {
-			return err
+			return 0, 0, "", err
 		}
 		if active != 1 {
-			return ErrWorkspaceInUse
+			return 0, 0, "", ErrWorkspaceInUse
 		}
-		chunk := workspaceCheckpointChunk{RunID: runID, ArchiveID: id, Sequence: count, Data: buffer[:n], CreatedAt: time.Now().UTC()}
+		chunk := workspaceCheckpointChunk{RunID: runID, ArchiveID: archiveID, Sequence: count, Data: buffer[:n], CreatedAt: time.Now().UTC()}
 		if err := r.DB.Create(&chunk).Error; err != nil {
-			return err
+			return 0, 0, "", err
 		}
-		hash.Write(buffer[:n])
+		digest.Write(buffer[:n])
 		count++
 		if err == io.ErrUnexpectedEOF {
 			break
 		}
 	}
-	if size == 0 {
-		return fmt.Errorf("empty workspace checkpoint")
-	}
-	err = r.DB.Transaction(func(tx *gorm.DB) error {
-		if _, err := r.takeRunForUpdate(tx, runID); err != nil {
-			return err
-		}
-		result := tx.Model(&WorkspaceCheckpoint{}).Where("run_id = ? AND lease_token = ? AND lease_until > ?", runID, token, time.Now().UTC()).Updates(map[string]any{
-			"archive": nil, "archive_id": id, "archive_bytes": size, "archive_chunks": count,
-			"sha256": hex.EncodeToString(hash.Sum(nil)), "saved_at": now, "error": "", "save_failures": 0, "save_retry_at": nil,
-		})
-		if result.Error != nil {
-			return result.Error
-		}
-		if result.RowsAffected != 1 {
-			return ErrWorkspaceInUse
-		}
-		return tx.Where("run_id = ? AND archive_id <> ?", runID, id).Delete(&workspaceCheckpointChunk{}).Error
-	})
-	published = err == nil
-	return err
+	return size, count, hex.EncodeToString(digest.Sum(nil)), nil
 }
 
 func (r *RunsRepo) OpenWorkspaceCheckpointArchive(runID string) (io.ReadCloser, error) {

internal/repo/workspace_checkpoint_stage_lease_guard_test.go

@@ -0,0 +1,85 @@
+package repo
+
+import (
+	"agentflow/internal/domain"
+	"bytes"
+	"errors"
+	"fmt"
+	"testing"
+	"time"
+)
+
+// leaseLosingReader serves data and runs lose once the first full chunk has
+// been handed out, i.e. while the stager is about to write the second row.
+type leaseLosingReader struct {
+	data   *bytes.Reader
+	served int
+	lose   func()
+	lost   bool
+}
+
+func (r *leaseLosingReader) Read(p []byte) (int, error) {
+	if !r.lost && r.served >= workspaceCheckpointChunkBytes {
+		r.lost = true
+		r.lose()
+	}
+	n, err := r.data.Read(p)
+	r.served += n
+	return n, err
+}
+
+// stageWorkspaceCheckpointChunks documents that it re-checks the lease before
+// every row. These guard that a writer whose lease is lost mid-stream stops
+// at the next row instead of only checking before the first one.
+func TestStageWorkspaceCheckpointChunksStopsWhenLeaseLostMidStream(t *testing.T) {
+	cases := []struct {
+		name string
+		lose func(s *Store, runID string) error
+	}{
+		{"lease expired", func(s *Store, runID string) error {
+			return s.DB.Model(&WorkspaceCheckpoint{}).Where("run_id = ?", runID).Update("lease_until", time.Now().UTC().Add(-time.Minute)).Error
+		}},
+		{"lease taken over", func(s *Store, runID string) error {
+			return s.DB.Model(&WorkspaceCheckpoint{}).Where("run_id = ?", runID).Update("lease_token", "other-writer").Error
+		}},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			s, err := OpenTest()
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer s.Close()
+			run := domain.Run{ID: fmt.Sprintf("checkpoint-stage-lease-%d", time.Now().UnixNano()), Status: "succeeded", Nodes: map[string]*domain.NodeRun{}}
+			if err := s.PutRun(run); err != nil {
+				t.Fatal(err)
+			}
+			token, err := s.BeginWorkspaceUse(run.ID, time.Now().UTC())
+			if err != nil {
+				t.Fatal(err)
+			}
+			source := &leaseLosingReader{
+				data: bytes.NewReader(bytes.Repeat([]byte("l"), 3*workspaceCheckpointChunkBytes)),
+				lose: func() {
+					if err := tc.lose(s, run.ID); err != nil {
+						t.Fatal(err)
+					}
+				},
+			}
+			_, _, _, err = s.stageWorkspaceCheckpointChunks(run.ID, token, "archive-lease-lost", source)
+			if !errors.Is(err, ErrWorkspaceInUse) {
+				t.Fatalf("writer kept staging after losing its lease: err=%v", err)
+			}
+			if !source.lost {
+				t.Fatal("lease was never lost; the test did not reach the second chunk")
+			}
+			var rows []workspaceCheckpointChunk
+			if err := s.DB.Where("run_id = ? AND archive_id = ?", run.ID, "archive-lease-lost").Order("sequence").Find(&rows).Error; err != nil {
+				t.Fatal(err)
+			}
+			if len(rows) != 1 || rows[0].Sequence != 0 {
+				t.Fatalf("only the row staged before the lease was lost may exist, got %d rows", len(rows))
+			}
+		})
+	}
+}

internal/repo/workspace_checkpoint_stage_test.go

@@ -0,0 +1,93 @@
+package repo
+
+import (
+	"agentflow/internal/domain"
+	"bytes"
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"testing"
+	"time"
+)
+
+func TestStageWorkspaceCheckpointChunks(t *testing.T) {
+	s, err := OpenTest()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer s.Close()
+	run := domain.Run{ID: fmt.Sprintf("checkpoint-stage-%d", time.Now().UnixNano()), Status: "succeeded", Nodes: map[string]*domain.NodeRun{}}
+	if err := s.PutRun(run); err != nil {
+		t.Fatal(err)
+	}
+	token, err := s.BeginWorkspaceUse(run.ID, time.Now().UTC())
+	if err != nil {
+		t.Fatal(err)
+	}
+	chunksOf := func(archiveID string) []workspaceCheckpointChunk {
+		var rows []workspaceCheckpointChunk
+		if err := s.DB.Where("run_id = ? AND archive_id = ?", run.ID, archiveID).Order("sequence").Find(&rows).Error; err != nil {
+			t.Fatal(err)
+		}
+		return rows
+	}
+
+	t.Run("multi chunk", func(t *testing.T) {
+		data := bytes.Repeat([]byte("stage"), (2*workspaceCheckpointChunkBytes+100)/5)
+		size, count, sum, err := s.stageWorkspaceCheckpointChunks(run.ID, token, "archive-multi", bytes.NewReader(data))
+		if err != nil {
+			t.Fatal(err)
+		}
+		want := sha256.Sum256(data)
+		if size != int64(len(data)) || count != 3 || sum != hex.EncodeToString(want[:]) {
+			t.Fatalf("size=%d count=%d sum=%s", size, count, sum)
+		}
+		rows := chunksOf("archive-multi")
+		if len(rows) != 3 {
+			t.Fatalf("rows=%d", len(rows))
+		}
+		var joined []byte
+		for i, row := range rows {
+			if row.Sequence != i || len(row.Data) > workspaceCheckpointChunkBytes {
+				t.Fatalf("row %d: sequence=%d len=%d", i, row.Sequence, len(row.Data))
+			}
+			joined = append(joined, row.Data...)
+		}
+		if !bytes.Equal(joined, data) {
+			t.Fatal("staged rows do not reassemble source")
+		}
+	})
+
+	t.Run("exact chunk boundary", func(t *testing.T) {
+		data := bytes.Repeat([]byte("x"), workspaceCheckpointChunkBytes)
+		size, count, _, err := s.stageWorkspaceCheckpointChunks(run.ID, token, "archive-exact", bytes.NewReader(data))
+		if err != nil || size != int64(len(data)) || count != 1 || len(chunksOf("archive-exact")) != 1 {
+			t.Fatalf("size=%d count=%d err=%v", size, count, err)
+		}
+	})
+
+	t.Run("empty source", func(t *testing.T) {
+		size, count, _, err := s.stageWorkspaceCheckpointChunks(run.ID, token, "archive-empty", bytes.NewReader(nil))
+		if err != nil || size != 0 || count != 0 || len(chunksOf("archive-empty")) != 0 {
+			t.Fatalf("size=%d count=%d err=%v", size, count, err)
+		}
+	})
+
+	t.Run("stale lease", func(t *testing.T) {
+		_, _, _, err := s.stageWorkspaceCheckpointChunks(run.ID, "stale-token", "archive-stale", bytes.NewReader([]byte("data")))
+		if !errors.Is(err, ErrWorkspaceInUse) {
+			t.Fatalf("err=%v", err)
+		}
+		if rows := chunksOf("archive-stale"); len(rows) != 0 {
+			t.Fatalf("stale writer staged %d rows", len(rows))
+		}
+	})
+
+	t.Run("read error", func(t *testing.T) {
+		_, _, _, err := s.stageWorkspaceCheckpointChunks(run.ID, token, "archive-fail", &checkpointFailingReader{remaining: workspaceCheckpointChunkBytes + 10})
+		if err == nil {
+			t.Fatal("read failure ignored")
+		}
+	})
+}