diff --git a/.gitignore b/.gitignore index 229d82812..3d2405047 100644 --- a/.gitignore +++ b/.gitignore @@ -69,6 +69,7 @@ CODEBUDDY.md .reasonix/ .mcp.json .zcodeignore +graphify-out/ # go core binary swanlab/bin/ diff --git a/core/internal/budget/budget.go b/core/internal/budget/budget.go new file mode 100644 index 000000000..072019658 --- /dev/null +++ b/core/internal/budget/budget.go @@ -0,0 +1,178 @@ +// Package budget 提供上传管线的会话级容量预算(条数与字节双限)。 +// +// 预算覆盖记录从入队到上传终结的整个生命周期:记录在组件间转移时不释放 +// 配额,只有同号覆盖或确认终结才归还。上限 <= 0 表示该维度不限制。 +package budget + +import ( + "context" + "errors" + "math" + "sync" +) + +// ErrTooLarge 表示单次申请量本身超过上限,等待也不会成功。 +var ErrTooLarge = errors.New("budget: request exceeds the configured limit") + +// Limits 是条数与字节上限,<=0 表示该维度不限制。 +type Limits struct { + Records int64 + Bytes int64 +} + +// Unlimited 报告两个维度是否都不限制。 +func (l Limits) Unlimited() bool { return l.Records <= 0 && l.Bytes <= 0 } + +// Budget 是并发安全的容量账本。零值不可用,应由 New 构造。 +type Budget struct { + mu sync.Mutex + limits Limits + usedRecords int64 + usedBytes int64 + // wait 在每次归还后关闭并重建,用于唤醒阻塞中的 Acquire。 + wait chan struct{} +} + +// New 创建容量账本;limits 的负值按不限制处理。 +func New(limits Limits) *Budget { + if limits.Records < 0 { + limits.Records = 0 + } + if limits.Bytes < 0 { + limits.Bytes = 0 + } + return &Budget{limits: limits, wait: make(chan struct{})} +} + +// Limits 返回构造时的上限。 +func (b *Budget) Limits() Limits { return b.limits } + +// Usage 返回当前占用的条数与字节数。 +func (b *Budget) Usage() (records, bytes int64) { + b.mu.Lock() + defer b.mu.Unlock() + return b.usedRecords, b.usedBytes +} + +// TryAcquire 尝试立即占用配额;额度不足或单次申请超过上限时返回 false, +// 且不改变任何状态。 +func (b *Budget) TryAcquire(records, bytes int64) bool { + checkAmount(records, bytes) + b.mu.Lock() + defer b.mu.Unlock() + if b.tooLarge(records, bytes) || !b.fitsLocked(records, bytes) { + return false + } + b.usedRecords += records + b.usedBytes += bytes + return true +} + +// Acquire 占用配额,额度不足时阻塞到有归还或 ctx 结束。 +// 单次申请量超过上限时立即返回 ErrTooLarge;不保证 FIFO。 +func (b *Budget) Acquire(ctx context.Context, records, bytes int64) error { + checkAmount(records, bytes) + if b.tooLarge(records, bytes) { + return ErrTooLarge + } + for { + b.mu.Lock() + if b.fitsLocked(records, bytes) { + b.usedRecords += records + b.usedBytes += bytes + b.mu.Unlock() + return nil + } + wait := b.wait + b.mu.Unlock() + + select { + case <-ctx.Done(): + return ctx.Err() + case <-wait: + } + } +} + +// Release 归还配额并唤醒等待者。 +// 超量归还属于调用约定错误,会 panic,避免账本静默失衡。 +func (b *Budget) Release(records, bytes int64) { + checkAmount(records, bytes) + b.mu.Lock() + defer b.mu.Unlock() + if records > b.usedRecords || bytes > b.usedBytes { + panic("budget: released more than reserved") + } + b.usedRecords -= records + b.usedBytes -= bytes + if records > 0 || bytes > 0 { + b.wakeLocked() + } +} + +// Adjust 按实际用量修正一次预留:正数补扣、负数退还。 +// 补扣可能使占用短暂超过上限;退款超量会 panic。 +func (b *Budget) Adjust(records, bytes int64) { + b.mu.Lock() + defer b.mu.Unlock() + if records < 0 && -records > b.usedRecords || bytes < 0 && -bytes > b.usedBytes { + panic("budget: adjusted more than reserved") + } + b.usedRecords += records + b.usedBytes += bytes + if records < 0 || bytes < 0 { + b.wakeLocked() + } +} + +// Full 报告任一有限维度是否已经用满。 +func (b *Budget) Full() bool { + b.mu.Lock() + defer b.mu.Unlock() + return b.limits.Records > 0 && b.usedRecords >= b.limits.Records || + b.limits.Bytes > 0 && b.usedBytes >= b.limits.Bytes +} + +// Ratio 返回占用率:条数与字节取较大值,无限维度不参与计算。 +// 两个维度都不限制时返回 0;经 Adjust 后可能短暂大于 1。 +func (b *Budget) Ratio() float64 { + b.mu.Lock() + defer b.mu.Unlock() + var ratio float64 + if b.limits.Records > 0 { + ratio = math.Max(ratio, float64(b.usedRecords)/float64(b.limits.Records)) + } + if b.limits.Bytes > 0 { + ratio = math.Max(ratio, float64(b.usedBytes)/float64(b.limits.Bytes)) + } + return ratio +} + +// wakeLocked 唤醒所有等待者;调用方需持有 mu。 +func (b *Budget) wakeLocked() { + close(b.wait) + b.wait = make(chan struct{}) +} + +// fitsLocked 判断当前占用加上本次申请是否仍在上限内;调用方需持有 mu。 +func (b *Budget) fitsLocked(records, bytes int64) bool { + if b.limits.Records > 0 && b.usedRecords+records > b.limits.Records { + return false + } + if b.limits.Bytes > 0 && b.usedBytes+bytes > b.limits.Bytes { + return false + } + return true +} + +// tooLarge 判断单次申请量是否已超过上限;limits 不可变,无需加锁。 +func (b *Budget) tooLarge(records, bytes int64) bool { + return b.limits.Records > 0 && records > b.limits.Records || + b.limits.Bytes > 0 && bytes > b.limits.Bytes +} + +func checkAmount(records, bytes int64) { + if records < 0 || bytes < 0 { + panic("budget: amount must be non-negative") + } +} diff --git a/core/internal/budget/budget_test.go b/core/internal/budget/budget_test.go new file mode 100644 index 000000000..3ae538e21 --- /dev/null +++ b/core/internal/budget/budget_test.go @@ -0,0 +1,323 @@ +package budget + +import ( + "context" + "errors" + "sync" + "testing" + "time" +) + +func TestNewNormalizesLimits(t *testing.T) { + b := New(Limits{Records: -1, Bytes: -5}) + if got := b.Limits(); got != (Limits{}) { + t.Fatalf("Limits() = %+v, want zero value", got) + } + if !b.Limits().Unlimited() { + t.Fatal("negative limits should be treated as unlimited") + } + if _, bytes := b.Usage(); bytes != 0 { + t.Fatalf("initial bytes = %d, want 0", bytes) + } +} + +func TestTryAcquire(t *testing.T) { + tests := []struct { + name string + limits Limits + first [2]int64 // records, bytes + second [2]int64 + wantOK bool + wantRec int64 + wantByt int64 + }{ + { + name: "unlimited always fits", + limits: Limits{}, + first: [2]int64{100, 1 << 20}, + second: [2]int64{100, 1 << 20}, + wantOK: true, + wantRec: 200, + wantByt: 2 << 20, + }, + { + name: "records limit", + limits: Limits{Records: 3}, + first: [2]int64{2, 0}, + second: [2]int64{2, 0}, + wantOK: false, + wantRec: 2, + }, + { + name: "bytes limit", + limits: Limits{Bytes: 10}, + first: [2]int64{1, 6}, + second: [2]int64{1, 5}, + wantOK: false, + wantRec: 1, + wantByt: 6, + }, + { + name: "exact fit", + limits: Limits{Records: 2, Bytes: 10}, + first: [2]int64{2, 10}, + second: [2]int64{1, 0}, + wantOK: false, + wantRec: 2, + wantByt: 10, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + b := New(tc.limits) + if !b.TryAcquire(tc.first[0], tc.first[1]) { + t.Fatalf("first TryAcquire(%d, %d) = false, want true", tc.first[0], tc.first[1]) + } + if got := b.TryAcquire(tc.second[0], tc.second[1]); got != tc.wantOK { + t.Fatalf("second TryAcquire(%d, %d) = %v, want %v", tc.second[0], tc.second[1], got, tc.wantOK) + } + records, bytes := b.Usage() + if records != tc.wantRec || bytes != tc.wantByt { + t.Fatalf("Usage() = (%d, %d), want (%d, %d)", records, bytes, tc.wantRec, tc.wantByt) + } + }) + } +} + +func TestTryAcquireTooLarge(t *testing.T) { + b := New(Limits{Records: 1}) + if b.TryAcquire(2, 0) { + t.Fatal("TryAcquire(2, 0) = true with limit 1") + } + if records, bytes := b.Usage(); records != 0 || bytes != 0 { + t.Fatalf("Usage() = (%d, %d), want (0, 0)", records, bytes) + } +} + +func TestAcquireWaitsForRelease(t *testing.T) { + b := New(Limits{Records: 1}) + if !b.TryAcquire(1, 0) { + t.Fatal("TryAcquire(1, 0) = false, want true") + } + + acquired := make(chan error, 1) + go func() { acquired <- b.Acquire(context.Background(), 1, 0) }() + + select { + case err := <-acquired: + t.Fatalf("Acquire returned before Release: %v", err) + case <-time.After(20 * time.Millisecond): + } + + b.Release(1, 0) + select { + case err := <-acquired: + if err != nil { + t.Fatalf("Acquire after Release = %v, want nil", err) + } + case <-time.After(time.Second): + t.Fatal("Acquire did not wake after Release") + } + if records, _ := b.Usage(); records != 1 { + t.Fatalf("Usage records = %d, want 1", records) + } +} + +func TestAcquireContextCanceled(t *testing.T) { + b := New(Limits{Records: 1}) + if !b.TryAcquire(1, 0) { + t.Fatal("TryAcquire(1, 0) = false, want true") + } + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + + err := b.Acquire(ctx, 1, 0) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("Acquire err = %v, want context.DeadlineExceeded", err) + } + if records, _ := b.Usage(); records != 1 { + t.Fatalf("Usage records = %d, want 1 (failed acquire must not change state)", records) + } +} + +func TestAcquireTooLarge(t *testing.T) { + b := New(Limits{Records: 2}) + err := b.Acquire(context.Background(), 3, 0) + if !errors.Is(err, ErrTooLarge) { + t.Fatalf("Acquire err = %v, want ErrTooLarge", err) + } + if records, _ := b.Usage(); records != 0 { + t.Fatalf("Usage records = %d, want 0", records) + } +} + +func TestReleasePanicsOnOverRelease(t *testing.T) { + tests := []struct { + name string + acquire [2]int64 + release [2]int64 + wantPanic bool + }{ + {name: "exact", acquire: [2]int64{2, 10}, release: [2]int64{2, 10}}, + {name: "over records", acquire: [2]int64{2, 10}, release: [2]int64{3, 10}, wantPanic: true}, + {name: "over bytes", acquire: [2]int64{2, 10}, release: [2]int64{2, 11}, wantPanic: true}, + {name: "negative", acquire: [2]int64{2, 10}, release: [2]int64{-1, 0}, wantPanic: true}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + b := New(Limits{}) + b.TryAcquire(tc.acquire[0], tc.acquire[1]) + defer func() { + got := recover() != nil + if got != tc.wantPanic { + t.Fatalf("panic = %v, want %v", got, tc.wantPanic) + } + }() + b.Release(tc.release[0], tc.release[1]) + }) + } +} + +func TestAdjust(t *testing.T) { + b := New(Limits{Records: 10, Bytes: 100}) + if !b.TryAcquire(2, 50) { + t.Fatal("TryAcquire(2, 50) = false, want true") + } + // 实际比估算小:退还差额。 + b.Adjust(1, -20) + if records, bytes := b.Usage(); records != 3 || bytes != 30 { + t.Fatalf("Usage after shrink = (%d, %d), want (3, 30)", records, bytes) + } + // 实际比估算大:补扣差额。 + b.Adjust(1, 20) + if records, bytes := b.Usage(); records != 4 || bytes != 50 { + t.Fatalf("Usage after grow = (%d, %d), want (4, 50)", records, bytes) + } +} + +func TestAdjustPanicsOnOverRefund(t *testing.T) { + b := New(Limits{}) + b.TryAcquire(1, 10) + defer func() { + if recover() == nil { + t.Fatal("Adjust with excessive refund did not panic") + } + }() + b.Adjust(0, -11) +} + +func TestWatermark(t *testing.T) { + b := New(Limits{Records: 4, Bytes: 100}) + if got := b.Ratio(); got != 0 { + t.Fatalf("empty Ratio = %v, want 0", got) + } + b.TryAcquire(2, 50) + if got := b.Ratio(); got != 0.5 { + t.Fatalf("Ratio = %v, want 0.5", got) + } + if b.Full() { + t.Fatal("Full() = true at 50%, want false") + } + b.TryAcquire(2, 50) + if got := b.Ratio(); got != 1 { + t.Fatalf("Ratio = %v, want 1", got) + } + if !b.Full() { + t.Fatal("Full() = false at 100%, want true") + } + + unlimited := New(Limits{}) + unlimited.TryAcquire(1<<20, 1<<30) + if got := unlimited.Ratio(); got != 0 { + t.Fatalf("unlimited Ratio = %v, want 0", got) + } + if unlimited.Full() { + t.Fatal("unlimited Full() = true, want false") + } +} + +// TestQuotaConservation 验证配额在「预留 → 组件间流转 → 覆盖/终结」全过程中守恒: +// 流转不释放,只有覆盖与终结才归还。 +func TestQuotaConservation(t *testing.T) { + tests := []struct { + name string + actions func(t *testing.T, b *Budget) + wantRec int64 + wantByt int64 + }{ + { + name: "transfer keeps quota", + actions: func(t *testing.T, b *Budget) { + // intake → upload channel → batcher: 只是换持有者。 + if !b.TryAcquire(2, 20) { + t.Fatal("TryAcquire = false") + } + }, + wantRec: 2, + wantByt: 20, + }, + { + name: "overwrite releases only the replaced record", + actions: func(t *testing.T, b *Budget) { + if !b.TryAcquire(2, 20) { + t.Fatal("TryAcquire = false") + } + // 同号覆盖:新记录占 1 条 10 字节,被覆盖的旧记录归还 1 条 10 字节。 + b.Release(1, 10) + if !b.TryAcquire(1, 10) { + t.Fatal("TryAcquire for replacement = false") + } + }, + wantRec: 2, + wantByt: 20, + }, + { + name: "confirmed releases", + actions: func(t *testing.T, b *Budget) { + if !b.TryAcquire(2, 20) { + t.Fatal("TryAcquire = false") + } + b.Release(1, 10) + b.Release(1, 10) + }, + wantRec: 0, + wantByt: 0, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + b := New(Limits{Records: 10, Bytes: 100}) + tc.actions(t, b) + records, bytes := b.Usage() + if records != tc.wantRec || bytes != tc.wantByt { + t.Fatalf("Usage() = (%d, %d), want (%d, %d)", records, bytes, tc.wantRec, tc.wantByt) + } + }) + } +} + +func TestConcurrentAcquireRelease(t *testing.T) { + const ( + workers = 8 + rounds = 200 + ) + b := New(Limits{Records: 4, Bytes: 40}) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < rounds; j++ { + if err := b.Acquire(context.Background(), 1, 10); err != nil { + t.Errorf("Acquire = %v", err) + return + } + b.Release(1, 10) + } + }() + } + wg.Wait() + if records, bytes := b.Usage(); records != 0 || bytes != 0 { + t.Fatalf("Usage after all workers = (%d, %d), want (0, 0)", records, bytes) + } +} diff --git a/core/internal/buffer/buffer.go b/core/internal/buffer/buffer.go new file mode 100644 index 000000000..18a6a1d18 --- /dev/null +++ b/core/internal/buffer/buffer.go @@ -0,0 +1,678 @@ +// Package buffer 实现会话入口分发器 intake:两阶段接收、容量预留与双通道路由。 +// +// Reserve 只做只读预检与容量预留;Commit 分配 storeSeq、写本地并分流到上传。 +// 本地 run-*.swanlab 是唯一事实源: +// - 数据类记录同时进入本地与上传; +// - Start(-1)/Finish(-2) 只落本地; +// - 崩溃 error_log 只上传; +// - SAVE_POLICY_END 先落本地、Close 时再上传。 +// +// 本地通道满时只做有界反压,调用方保留数据后重试;只有真实写失败才降级为 +// 上传兜底。暂停投递期间保留未确认落盘的记录,写盘失败时重新投入上传。 +// skip_store 是启动期配置,运行期磁盘写满只置 storeDegraded,不会修改配置。 +package buffer + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "time" + + "google.golang.org/protobuf/proto" + + "github.com/swanhubx/swanlab/core/internal/budget" + "github.com/swanhubx/swanlab/core/internal/pkg/console" + "github.com/swanhubx/swanlab/core/internal/pkg/record" + "github.com/swanhubx/swanlab/core/internal/store" + recordv1 "github.com/swanhubx/swanlab/core/proto/swanlab/record/v1" + savev1 "github.com/swanhubx/swanlab/core/proto/swanlab/save/v1" +) + +const ( + // DefaultQueueWait 是本地通道、上传预算或上传槽位不足时的有界等待上限。 + DefaultQueueWait = 2 * time.Second + // DefaultUploadQueue 是上传通道的默认容量,与单批上限对齐。 + DefaultUploadQueue = 10_000 + // defaultStorePollInterval 是等待本地或上传通道腾出空间时的轮询间隔。 + defaultStorePollInterval = 10 * time.Millisecond +) + +var ( + // ErrClosed 表示 intake 入口已经封闭。 + ErrClosed = errors.New("buffer: intake is closed") + // ErrOverloaded 表示容量在等待上限内没有释放,调用方保留数据后可重试。 + ErrOverloaded = errors.New("buffer: capacity exhausted") + // ErrBatchTooLarge 表示单批超过通道容量,重试不会成功。 + ErrBatchTooLarge = errors.New("buffer: batch exceeds queue capacity") + // ErrOverReserved 表示本次提交超出 Reserve 预留的配额,属于调用方错误。 + ErrOverReserved = errors.New("buffer: commit exceeds reservation") + // errQueueWait 标记本地通道等待超时;与调用方 ctx 取消区分开。 + errQueueWait = errors.New("buffer: local queue wait exceeded") +) + +// Options 是 Intake 的构造参数。 +type Options struct { + // SkipStore 是启动期配置,创建后只读;为 true 时不写本地,Store 必须为 nil。 + SkipStore bool + // Store 是 AsyncWriter 的写入入口,容量应不小于单批最大记录数。 + Store chan<- store.Request + // UploadQueue 是上传通道容量,应不小于单批最大记录数;<=0 取 DefaultUploadQueue。 + UploadQueue int + // Budget 是会话级上传预算。 + Budget *budget.Budget + // Degraded 是运行期 storeDegraded 标志:写/Flush/序列化失败置位, + // 容量不足不置位;不会回写配置。为 nil 时内部自建(测试用)。 + Degraded *atomic.Bool + // Flushed 返回已完整落盘的 storeSeq;!SkipStore 时必填。 + // 用于确认暂停期间的记录已安全落盘,不得回调 Intake。 + Flushed func() uint64 + // QueueWait 是容量不足时的有界等待上限,<=0 取 DefaultQueueWait。 + QueueWait time.Duration + // StorePollInterval 是本地通道等待的轮询间隔,<=0 取 10ms。 + StorePollInterval time.Duration +} + +// Demand 是一次提交对两条通道的容量需求。 +type Demand struct { + // Store 是需要写入本地通道的记录数;SkipStore 时忽略。 + Store int + // Upload 是需要立即进入上传通道的记录数(不含 Deferred)。 + Upload int + // Deferred 是 SAVE_POLICY_END 记录数:先落本地,Close 时上传; + // 其配额从 Reserve 起一直占用到上传终结。 + Deferred int + // Bytes 是 Upload 记录的序列化字节数估算。 + Bytes int64 + // DeferredBytes 是 Deferred 记录的序列化字节数估算;暂停时只按它申请配额。 + // Commit 会按实际大小修正补扣,可能短暂超出上限。 + DeferredBytes int64 +} + +// Ticket 是 Reserve 返回的预留凭证,由 Commit 消费或 Release 归还。 +type Ticket struct { + store int + upload int // 已占用的上传预算条数,含延迟记录 + slots int // 已预留的上传通道槽位,仅覆盖立即投递的记录 + bytes int64 + paused bool +} + +// Intake 是单个会话的入口分发器,并发安全。 +type Intake struct { + skipStore bool + storeOut chan<- store.Request + upload chan *recordv1.Record + budget *budget.Budget + degraded *atomic.Bool + flushed func() uint64 + wait time.Duration + poll time.Duration + uploadCap int + + mu sync.Mutex + closeMu sync.Mutex + closeOnce sync.Once + nextSeq uint64 + pendingStore int + pendingUpload int + pending []*recordv1.Record + unconfirmed []unconfirmedRec + recovery []*recordv1.Record + pauseSeq uint64 // 0 表示未暂停;非零为首个未投递记录的 storeSeq + closed bool +} + +// New 校验参数并创建 intake。 +func New(opts Options) (*Intake, error) { + if opts.Budget == nil { + return nil, errors.New("buffer: nil budget") + } + if !opts.SkipStore && opts.Store == nil { + return nil, errors.New("buffer: nil store channel") + } + if !opts.SkipStore && opts.Flushed == nil { + return nil, errors.New("buffer: nil flushed callback") + } + if opts.QueueWait <= 0 { + opts.QueueWait = DefaultQueueWait + } + if opts.StorePollInterval <= 0 { + opts.StorePollInterval = defaultStorePollInterval + } + degraded := opts.Degraded + if degraded == nil { + degraded = &atomic.Bool{} + } + queue := opts.UploadQueue + if queue <= 0 { + queue = DefaultUploadQueue + } + upload := make(chan *recordv1.Record, queue) + t := &Intake{ + skipStore: opts.SkipStore, + storeOut: opts.Store, + upload: upload, + budget: opts.Budget, + degraded: degraded, + flushed: opts.Flushed, + wait: opts.QueueWait, + poll: opts.StorePollInterval, + uploadCap: cap(upload), + } + return t, nil +} + +// UploadOut 返回上传通道;intake 在 Close 时封闭它,消费方读尽缓冲区后退出。 +func (t *Intake) UploadOut() <-chan *recordv1.Record { return t.upload } + +// Reserve 预留本地通道、上传预算与上传槽位,不修改记录状态。 +// +// 预留失败不产生任何副作用。本地通道满时在 QueueWait 内有界等待,仍失败 +// 返回 ErrOverloaded,由调用方保留数据后重试。 +// +// 本地副本完整但上传预算不足时,立即记录转为暂停投递、只落本地;延迟记录 +// 必须从接收到 Close 一直持有配额,占不到时立即返回 ErrOverloaded,避免 +// 暂停期间暂存无界增长。没有完整本地副本时禁止暂停投递,所有记录有界 +// 阻塞等待上传预算,失败返回 ErrOverloaded。 +func (t *Intake) Reserve(ctx context.Context, d Demand) (Ticket, error) { + if d.Store < 0 || d.Upload < 0 || d.Deferred < 0 || d.Bytes < 0 || d.DeferredBytes < 0 { + return Ticket{}, fmt.Errorf("buffer: invalid demand store=%d upload=%d deferred=%d bytes=%d deferredBytes=%d", + d.Store, d.Upload, d.Deferred, d.Bytes, d.DeferredBytes) + } + + t.mu.Lock() + if t.closed { + t.mu.Unlock() + return Ticket{}, ErrClosed + } + paused := t.pauseSeq != 0 + if t.degraded.Load() { + t.collectUnconfirmedLocked() + } + t.trimUnconfirmedLocked() + t.mu.Unlock() + + local := !t.skipStore && !t.degraded.Load() + tk, err := t.reserveStoreSide(ctx, d, local) + if err != nil { + return Ticket{}, err + } + if d.Upload == 0 && d.Deferred == 0 { + return tk, nil + } + return t.reserveUploadSide(ctx, d, local, paused, tk) +} + +// reserveStoreSide 预留本地通道;降级或 skip_store 时归还本地预留。 +func (t *Intake) reserveStoreSide(ctx context.Context, d Demand, local bool) (Ticket, error) { + var tk Ticket + if local && d.Store > 0 { + err := t.reserveStore(ctx, d.Store) + switch { + case err == nil: + tk.store = d.Store + case errors.Is(err, ErrBatchTooLarge): + return Ticket{}, err + case errors.Is(err, errQueueWait): + // 本地是唯一事实源:容量不足时反压调用方重试,绝不丢弃/降级。 + return Ticket{}, overload(err) + case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): + return Ticket{}, err + default: + return Ticket{}, overload(err) + } + } + if !local { + t.releaseStore(tk.store) + tk.store = 0 + } + return tk, nil +} + +// reserveUploadSide 预留上传预算与通道槽位;本地副本完整且预算满时转为暂停投递。 +func (t *Intake) reserveUploadSide(ctx context.Context, d Demand, local, paused bool, tk Ticket) (Ticket, error) { + total := d.Upload + d.Deferred + totalBytes := d.Bytes + d.DeferredBytes + switch { + case !local: + // 无完整本地副本:立即与延迟记录都必须占用配额,禁止暂停投递。 + actx, cancel := context.WithTimeoutCause(ctx, t.wait, errQueueWait) + err := t.budget.Acquire(actx, int64(total), totalBytes) + cause := context.Cause(actx) + cancel() + if err != nil { + t.releaseStore(tk.store) + if errors.Is(cause, errQueueWait) { + return Ticket{}, overload(err) + } + return Ticket{}, err + } + tk.upload, tk.bytes = total, totalBytes + case paused: + // 已在暂停:立即记录依赖本地副本,延迟记录必须全程持有配额。 + if !t.holdDeferred(d, &tk) { + t.releaseStore(tk.store) + return Ticket{}, overload(nil) + } + tk.paused = true + return tk, nil + case !t.budget.TryAcquire(int64(total), totalBytes): + // 预算已满:只标记票据暂停,由 Commit 在真正提交时发布全局暂停状态。 + // 票据被放弃或提交失败时不能遗留暂停,否则后续批次会绕过预算获取。 + if !t.holdDeferred(d, &tk) { + t.releaseStore(tk.store) + return Ticket{}, overload(nil) + } + tk.paused = true + return tk, nil + default: + tk.upload, tk.bytes = total, totalBytes + } + + // 上传通道槽位只覆盖立即投递的记录,延迟记录在 Close 时才入队。 + if err := t.reserveUpload(ctx, d.Upload); err != nil { + t.budget.Release(int64(tk.upload), tk.bytes) + t.releaseStore(tk.store) + if errors.Is(err, ErrBatchTooLarge) || + errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return Ticket{}, err + } + return Ticket{}, overload(err) + } + tk.slots = d.Upload + return tk, nil +} + +// holdDeferred 为延迟记录占用上传配额。预算不足时由调用方反压, +// 不能先暂存到 Close 再补扣,否则 pending 会脱离预算无界增长。 +func (t *Intake) holdDeferred(d Demand, tk *Ticket) bool { + if d.Deferred == 0 { + return true + } + // 只按延迟记录自身的字节申请,避免大立即记录拖累小 END 保存。 + if !t.budget.TryAcquire(int64(d.Deferred), d.DeferredBytes) { + return false + } + tk.upload, tk.bytes = d.Deferred, d.DeferredBytes + return true +} + +// Commit 分配 storeSeq、写本地并分流到上传。票据在返回后即被消费, +// 未使用的预留会自动归还。 +// +// 本地已降级时跳过本地写入、继续上传;暂停投递的批次若本地副本已不可用, +// 会改走上传(可能短暂超出预算)以免丢记录。 +func (t *Intake) Commit(tk Ticket, records []*recordv1.Record) error { + if len(records) == 0 { + t.Release(tk) + return nil + } + for _, rec := range records { + if rec == nil { + t.Release(tk) + return errors.New("buffer: nil record") + } + } + + immediate, deferred := routeRecords(records) + uploadable := len(immediate) + len(deferred) + payloads, wantStore := t.encodeStore(records) + + t.mu.Lock() + if t.closed { + t.mu.Unlock() + t.Release(tk) + return ErrClosed + } + if wantStore && len(records) > tk.store { + t.mu.Unlock() + t.Release(tk) + return fmt.Errorf("%w: store %d > reserved %d", ErrOverReserved, len(records), tk.store) + } + // 票据自身的暂停标记在 ResumeUpload 之后仍然生效。 + paused := t.pauseSeq != 0 || tk.paused + if paused { + // 暂停批次只需为延迟记录持有配额,立即记录只落本地。 + if len(deferred) > tk.upload { + t.mu.Unlock() + t.Release(tk) + return fmt.Errorf("%w: deferred %d > reserved %d", ErrOverReserved, len(deferred), tk.upload) + } + } else if uploadable > tk.upload { + t.mu.Unlock() + t.Release(tk) + return fmt.Errorf("%w: upload %d > reserved %d", ErrOverReserved, uploadable, tk.upload) + } else if len(immediate) > tk.slots { + t.mu.Unlock() + t.Release(tk) + return fmt.Errorf("%w: immediate %d > slots %d", ErrOverReserved, len(immediate), tk.slots) + } + if !wantStore && tk.store > 0 { + t.releaseStoreLocked(tk.store) + tk.store = 0 + } + t.trimUnconfirmedLocked() + + // 本地通道:预留空间在握,发送不会阻塞。 + seqOf := make([]uint64, len(records)) + t.sendStoreLocked(records, payloads, wantStore, seqOf) + if tk.store > len(records) { + t.releaseStoreLocked(tk.store - len(records)) + } + + storeOK := wantStore && !t.degraded.Load() + if !storeOK && !t.skipStore { + // 本地已不可用:把此前暂停、尚未确认落盘的记录转入回收队列。 + t.collectUnconfirmedLocked() + } + pausedNow := false + switch { + case paused && storeOK: + pausedNow = t.pauseLocked(records, immediate, deferred, tk, seqOf) + case !storeOK && !t.skipStore: + // 本批没有可靠本地副本:并入回收队列按序补传,不能越过更早的记录。 + t.recoveryUploadLocked(records, immediate, deferred, tk) + default: + t.uploadLocked(records, immediate, deferred, tk) + } + t.mu.Unlock() + + if pausedNow { + console.Warning("upload budget is full, pausing upload delivery; records will stay in the local store") + } + return nil +} + +// routeRecords 按记录类型拆出立即上传与延迟上传(SAVE_POLICY_END)的下标。 +func routeRecords(records []*recordv1.Record) (immediate, deferred []int) { + immediate = make([]int, 0, len(records)) + deferred = make([]int, 0, len(records)) + for i, rec := range records { + switch record.TypeOf(rec) { + case record.TypeStart, record.TypeFinish: + // 仅落本地。 + case record.TypeSave: + if rec.GetSave().GetPolicy() == savev1.SavePolicy_SAVE_POLICY_END { + deferred = append(deferred, i) + continue + } + immediate = append(immediate, i) + default: + immediate = append(immediate, i) + } + } + return immediate, deferred +} + +// encodeStore 预先序列化落盘 payload;失败则降级并继续上传。 +func (t *Intake) encodeStore(records []*recordv1.Record) ([][]byte, bool) { + if t.skipStore || t.degraded.Load() { + return nil, false + } + payloads := make([][]byte, len(records)) + for i, rec := range records { + payload, err := proto.Marshal(rec) + if err != nil { + t.degrade(fmt.Errorf("marshal record %d for local store: %w", rec.GetNum(), err)) + return nil, false + } + payloads[i] = payload + } + return payloads, true +} + +// sendStoreLocked 把记录写入本地通道并分配 storeSeq;调用方需持有 mu。 +func (t *Intake) sendStoreLocked(records []*recordv1.Record, payloads [][]byte, wantStore bool, seqOf []uint64) { + if !wantStore { + return + } + for i := range records { + t.nextSeq++ + seqOf[i] = t.nextSeq + t.storeOut <- store.Request{StoreSeq: t.nextSeq, Payload: payloads[i]} + t.pendingStore-- + } +} + +// uploadLocked 按实际大小修正预算并投入上传通道;需持有 mu。 +// 上传槽位已在 Reserve 预留,发送不会阻塞。 +func (t *Intake) uploadLocked(records []*recordv1.Record, immediate, deferred []int, tk Ticket) { + actual := int64(0) + for _, i := range immediate { + actual += record.Size(records[i]) + } + for _, i := range deferred { + actual += record.Size(records[i]) + } + t.budget.Adjust(int64(len(immediate)+len(deferred)-tk.upload), actual-tk.bytes) + for _, i := range immediate { + t.upload <- records[i] + } + for _, i := range deferred { + t.pending = append(t.pending, records[i]) + } + // 已发出的记录由通道长度计入容量,预留槽位整体归还。 + t.releaseUploadLocked(tk.slots) +} + +// CommitUploadOnly 提交只上传、不落盘的记录(崩溃 error_log)。 +// 没有本地副本,因此在预算内阻塞等待,失败返回错误由调用方重试。 +func (t *Intake) CommitUploadOnly(ctx context.Context, rec *recordv1.Record) error { + if rec == nil { + return errors.New("buffer: nil record") + } + size := record.Size(rec) + if err := t.budget.Acquire(ctx, 1, size); err != nil { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return err + } + return overload(err) + } + if err := t.reserveUpload(ctx, 1); err != nil { + t.budget.Release(1, size) + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return err + } + return overload(err) + } + + t.mu.Lock() + defer t.mu.Unlock() + if t.closed { + t.budget.Release(1, size) + t.releaseUploadLocked(1) + return ErrClosed + } + t.upload <- rec + t.releaseUploadLocked(1) + return nil +} + +// Release 归还一张未被 Commit 消费的票据;可安全用于零值票据。 +func (t *Intake) Release(tk Ticket) { + if tk.store == 0 && tk.upload == 0 && tk.slots == 0 { + return + } + t.mu.Lock() + defer t.mu.Unlock() + t.releaseStoreLocked(tk.store) + if tk.upload > 0 { + t.budget.Release(int64(tk.upload), tk.bytes) + } + t.releaseUploadLocked(tk.slots) +} + +// Close 先补投回收队列与暂存的 SAVE_POLICY_END 记录,再关闭上传通道;可重复调用。 +// 调用方应在本方法之前关闭本地 store,并在其失败时先调用 DrainRecovery。 +func (t *Intake) Close(ctx context.Context) error { + t.closeMu.Lock() + defer t.closeMu.Unlock() + + t.mu.Lock() + t.closed = true + if t.degraded.Load() { + t.collectUnconfirmedLocked() + } else { + t.unconfirmed = nil + } + pending := t.pending + t.pending = nil + recovery := t.recovery + t.recovery = nil + t.mu.Unlock() + + for i := range recovery { + size := record.Size(recovery[i]) + t.budget.Adjust(1, size) + select { + case t.upload <- recovery[i]: + case <-ctx.Done(): + t.budget.Adjust(-1, -size) + t.restoreRecovery(recovery[i:]) + // 延迟记录尚未开始投递,必须随本次失败一起放回。 + t.restorePending(pending) + return ctx.Err() + } + } + for i := range pending { + select { + case t.upload <- pending[i]: + case <-ctx.Done(): + t.restorePending(pending[i:]) + return ctx.Err() + } + } + + t.closeOnce.Do(func() { close(t.upload) }) + return nil +} + +// reserveStore 在有界等待内为 n 条记录预留本地通道空间。 +// 容量不足只返回可重试错误;队列满时按 poll 间隔重试。 +func (t *Intake) reserveStore(ctx context.Context, n int) error { + if n > cap(t.storeOut) { + return fmt.Errorf("%w: batch %d > capacity %d", ErrBatchTooLarge, n, cap(t.storeOut)) + } + return t.reserveSlots(ctx, n, t.tryReserveStore) +} + +// reserveUpload 在有界等待内为 n 条记录预留上传通道槽位。 +func (t *Intake) reserveUpload(ctx context.Context, n int) error { + if n <= 0 { + return nil + } + if n > t.uploadCap { + return fmt.Errorf("%w: upload batch %d > capacity %d", ErrBatchTooLarge, n, t.uploadCap) + } + return t.reserveSlots(ctx, n, t.tryReserveUpload) +} + +// reserveSlots 在有界等待内反复尝试占用 n 个槽位。 +func (t *Intake) reserveSlots(ctx context.Context, n int, try func(int) bool) error { + if try(n) { + return nil + } + actx, cancel := context.WithTimeoutCause(ctx, t.wait, errQueueWait) + defer cancel() + ticker := time.NewTicker(t.poll) + defer ticker.Stop() + for { + select { + case <-actx.Done(): + return context.Cause(actx) + case <-ticker.C: + } + if try(n) { + return nil + } + } +} + +// tryReserveStore 在空间足够时原子占用 n 个位置。 +func (t *Intake) tryReserveStore(n int) bool { + t.mu.Lock() + defer t.mu.Unlock() + if cap(t.storeOut)-len(t.storeOut)-t.pendingStore < n { + return false + } + t.pendingStore += n + return true +} + +// releaseStore 归还本地通道预留;调用方不得持有 mu。 +func (t *Intake) releaseStore(n int) { + if n <= 0 { + return + } + t.mu.Lock() + defer t.mu.Unlock() + t.releaseStoreLocked(n) +} + +// releaseStoreLocked 归还本地通道预留;调用方需持有 mu。 +func (t *Intake) releaseStoreLocked(n int) { + if n <= 0 { + return + } + if n > t.pendingStore { + panic("buffer: released more local slots than reserved") + } + t.pendingStore -= n +} + +func (t *Intake) tryReserveUpload(n int) bool { + t.mu.Lock() + defer t.mu.Unlock() + return t.tryReserveUploadLocked(n) +} + +// tryReserveUploadLocked 在槽位足够时原子占用 n 个位置;调用方需持有 mu。 +func (t *Intake) tryReserveUploadLocked(n int) bool { + if n <= 0 { + return true + } + if t.uploadCap-len(t.upload)-t.pendingUpload < n { + return false + } + t.pendingUpload += n + return true +} + +// releaseUploadLocked 归还上传通道槽位;调用方需持有 mu。 +func (t *Intake) releaseUploadLocked(n int) { + if n <= 0 { + return + } + if n > t.pendingUpload { + panic("buffer: released more upload slots than reserved") + } + t.pendingUpload -= n +} + +// restorePending 把未发出的延迟记录放回队头,保持原有顺序。 +func (t *Intake) restorePending(items []*recordv1.Record) { + if len(items) == 0 { + return + } + t.mu.Lock() + rest := append([]*recordv1.Record(nil), items...) + t.pending = append(rest, t.pending...) + t.mu.Unlock() +} + +// degrade 置为本地存储降级状态并告警一次,只在真实写入失败时调用。 +func (t *Intake) degrade(err error) { + if t.degraded.CompareAndSwap(false, true) { + console.Error("local store is unavailable, continuing without a local copy:", err) + } +} + +// overload 包装为可重试的容量耗尽错误,保留底层原因。 +func overload(err error) error { + if err == nil { + return ErrOverloaded + } + return fmt.Errorf("%w: %v", ErrOverloaded, err) +} diff --git a/core/internal/buffer/buffer_test.go b/core/internal/buffer/buffer_test.go new file mode 100644 index 000000000..93706b144 --- /dev/null +++ b/core/internal/buffer/buffer_test.go @@ -0,0 +1,1328 @@ +package buffer + +import ( + "context" + "errors" + "fmt" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/swanhubx/swanlab/core/internal/budget" + "github.com/swanhubx/swanlab/core/internal/pkg/record" + "github.com/swanhubx/swanlab/core/internal/store" + columnv1 "github.com/swanhubx/swanlab/core/proto/swanlab/metric/column/v1" + datav1 "github.com/swanhubx/swanlab/core/proto/swanlab/metric/data/v1" + recordv1 "github.com/swanhubx/swanlab/core/proto/swanlab/record/v1" + runv1 "github.com/swanhubx/swanlab/core/proto/swanlab/run/v1" + savev1 "github.com/swanhubx/swanlab/core/proto/swanlab/save/v1" + terminalv1 "github.com/swanhubx/swanlab/core/proto/swanlab/terminal/v1" +) + +func scalarRecord(num int64) *recordv1.Record { + return &recordv1.Record{ + Num: num, + RecordType: &recordv1.Record_Scalar{Scalar: &datav1.ScalarRecord{ + Key: "train/loss", + Step: num, + Type: columnv1.ColumnType_COLUMN_TYPE_SCALAR, + Value: &datav1.ScalarValue{Number: float64(num)}, + }}, + } +} + +func startRecord() *recordv1.Record { + return &recordv1.Record{Num: record.NumStart, RecordType: &recordv1.Record_Start{Start: &runv1.StartRecord{}}} +} + +func finishRecord() *recordv1.Record { + return &recordv1.Record{Num: record.NumFinish, RecordType: &recordv1.Record_Finish{Finish: &runv1.FinishRecord{}}} +} + +func logRecord(num int64) *recordv1.Record { + return &recordv1.Record{Num: num, RecordType: &recordv1.Record_Log{Log: &terminalv1.LogRecord{Line: "line", Epoch: num}}} +} + +func saveRecord(num int64, policy savev1.SavePolicy) *recordv1.Record { + return &recordv1.Record{ + Num: num, + RecordType: &recordv1.Record_Save{Save: &savev1.SaveRecord{ + Name: fmt.Sprintf("file-%d", num), + Policy: policy, + Type: savev1.SaveType_SAVE_TYPE_CUSTOM, + }}, + } +} + +func demandFor(recs []*recordv1.Record) Demand { + var d Demand + for _, rec := range recs { + d.Store++ + size := record.Size(rec) + switch record.TypeOf(rec) { + case record.TypeStart, record.TypeFinish: + case record.TypeSave: + if rec.GetSave().GetPolicy() == savev1.SavePolicy_SAVE_POLICY_END { + d.Deferred++ + d.DeferredBytes += size + continue + } + d.Upload++ + d.Bytes += size + default: + d.Upload++ + d.Bytes += size + } + } + return d +} + +type storeSink struct { + ch chan store.Request + mu sync.Mutex + got []store.Request + notify chan struct{} + done chan struct{} +} + +func newStoreSink(size int) *storeSink { + s := &storeSink{ch: make(chan store.Request, size), notify: make(chan struct{}, 1), done: make(chan struct{})} + go func() { + defer close(s.done) + for req := range s.ch { + s.mu.Lock() + s.got = append(s.got, req) + s.mu.Unlock() + select { + case s.notify <- struct{}{}: + default: + } + } + }() + return s +} + +func (s *storeSink) stop() { + close(s.ch) + <-s.done +} + +func (s *storeSink) snapshot() []store.Request { + s.mu.Lock() + defer s.mu.Unlock() + return append([]store.Request(nil), s.got...) +} + +func newIntake(t *testing.T, opts Options) *Intake { + t.Helper() + if !opts.SkipStore && opts.Flushed == nil { + // 默认视为全部已刷盘;需要验证未确认记录时由用例自行注入。 + opts.Flushed = func() uint64 { return ^uint64(0) } + } + in, err := New(opts) + if err != nil { + t.Fatalf("New: %v", err) + } + return in +} + +// mustReserve 按记录数推导预留并返回票据。 +func mustReserve(t *testing.T, in *Intake, recs ...*recordv1.Record) Ticket { + t.Helper() + tk, err := in.Reserve(context.Background(), demandFor(recs)) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + return tk +} + +// mustCommit 按记录数推导预留并提交一批记录。 +func mustCommit(t *testing.T, in *Intake, recs ...*recordv1.Record) { + t.Helper() + tk := mustReserve(t, in, recs...) + if err := in.Commit(tk, recs); err != nil { + t.Fatalf("Commit: %v", err) + } +} + +func drainUpload(in *Intake) []*recordv1.Record { + var out []*recordv1.Record + for { + select { + case rec, ok := <-in.UploadOut(): + if !ok { + return out + } + out = append(out, rec) + default: + return out + } + } +} + +func waitStore(t *testing.T, sink *storeSink, want int) []store.Request { + t.Helper() + deadline := time.After(time.Second) + for len(sink.snapshot()) < want { + select { + case <-sink.notify: + case <-deadline: + t.Fatalf("store received %d records, want %d", len(sink.snapshot()), want) + } + } + return sink.snapshot() +} + +func TestNewValidatesOptions(t *testing.T) { + if _, err := New(Options{}); err == nil { + t.Fatal("New without budget returned nil error") + } + b := budget.New(budget.Limits{}) + if _, err := New(Options{Budget: b}); err == nil { + t.Fatal("New without store channel returned nil error") + } + if _, err := New(Options{Budget: b, Store: make(chan store.Request, 1)}); err == nil { + t.Fatal("New without flushed callback returned nil error") + } + if _, err := New(Options{Budget: b, SkipStore: true}); err != nil { + t.Fatalf("New with skip_store: %v", err) + } +} + +func TestCommitEmptyAndNil(t *testing.T) { + sink := newStoreSink(4) + defer sink.stop() + b := budget.New(budget.Limits{Records: 4}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 4}) + + tk, err := in.Reserve(context.Background(), Demand{Store: 1, Upload: 1, Bytes: 10}) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + if err = in.Commit(tk, nil); err != nil { + t.Fatalf("Commit empty: %v", err) + } + if records, bytes := b.Usage(); records != 0 || bytes != 0 { + t.Fatalf("budget usage = (%d, %d), want (0, 0)", records, bytes) + } + + tk, err = in.Reserve(context.Background(), Demand{Store: 1, Upload: 1, Bytes: 10}) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + if err = in.Commit(tk, []*recordv1.Record{nil}); err == nil { + t.Fatal("Commit with nil record returned nil error") + } + if records, bytes := b.Usage(); records != 0 || bytes != 0 { + t.Fatalf("budget usage after nil record = (%d, %d), want (0, 0)", records, bytes) + } +} + +// TestPausedEndSavesHoldBudget 覆盖暂停期间连续提交 END 保存:延迟记录必须 +// 从 Reserve 起占用上传预算,预算不足时反压,不能先暂存到 Close 再补扣。 +func TestPausedEndSavesHoldBudget(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{Records: 1}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 4}) + + first := []*recordv1.Record{scalarRecord(1)} + mustCommit(t, in, first...) + if got := len(drainUpload(in)); got != 1 { + t.Fatalf("upload got %d, want 1", got) + } + // 预算已满 → 下一条立即记录进入暂停并发布暂停状态。 + mustCommit(t, in, scalarRecord(2)) + if paused, seq := in.PauseState(); !paused || seq != 2 { + t.Fatalf("PauseState = (%v, %d), want (true, 2)", paused, seq) + } + + // 暂停期间预算已满:END 保存必须被反压,不得无配额暂存。 + save := []*recordv1.Record{saveRecord(3, savev1.SavePolicy_SAVE_POLICY_END)} + for i := 0; i < 2; i++ { + if _, err := in.Reserve(context.Background(), demandFor(save)); !errors.Is(err, ErrOverloaded) { + t.Fatalf("Reserve END save #%d err = %v, want ErrOverloaded", i+1, err) + } + } + if records, _ := b.Usage(); records != 1 { + t.Fatalf("budget usage = %d, want 1 (rejected saves must not be charged)", records) + } + if got := waitStore(t, sink, 2); len(got) != 2 { + t.Fatalf("store got %d, want 2", len(got)) + } + + // 模拟 uploader 确认后归还配额:END 保存被接收,并从 Reserve 起持续占用配额。 + records, bytes := b.Usage() + b.Release(records, bytes) + mustCommit(t, in, save...) + if records, _ := b.Usage(); records != 1 { + t.Fatalf("budget usage = %d, want 1 (deferred save charged at reserve)", records) + } + second := []*recordv1.Record{saveRecord(4, savev1.SavePolicy_SAVE_POLICY_END)} + if _, err := in.Reserve(context.Background(), demandFor(second)); !errors.Is(err, ErrOverloaded) { + t.Fatalf("second pending save err = %v, want ErrOverloaded", err) + } + + if err := in.Close(context.Background()); err != nil { + t.Fatalf("Close: %v", err) + } + got := drainUpload(in) + if len(got) != 1 || got[0].GetNum() != 3 { + t.Fatalf("upload after close = %v, want record num 3", recordNums(got)) + } + if records, _ := b.Usage(); records != 1 { + t.Fatalf("budget usage = %d, want 1 until the uploader confirms", records) + } +} + +// TestPausedMixedBatchHoldsOnlyDeferredBytes 覆盖“立即记录字节超限、END 保存 +// 未超限”的混合批次:暂停只需为延迟记录申请自身的字节配额。 +func TestPausedMixedBatchHoldsOnlyDeferredBytes(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + save := saveRecord(1, savev1.SavePolicy_SAVE_POLICY_END) + immediate := scalarRecord(2) + immediate.GetScalar().Key = strings.Repeat("k", 128) + if record.Size(immediate) <= record.Size(save) { + t.Fatalf("immediate bytes %d <= deferred bytes %d", record.Size(immediate), record.Size(save)) + } + b := budget.New(budget.Limits{Bytes: record.Size(save) + 8}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 4}) + + recs := []*recordv1.Record{immediate, save} + tk, err := in.Reserve(context.Background(), demandFor(recs)) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + if !tk.paused { + t.Fatal("mixed batch should pause upload when the immediate record exceeds the byte limit") + } + if err = in.Commit(tk, recs); err != nil { + t.Fatalf("Commit: %v", err) + } + if _, bytes := b.Usage(); bytes != record.Size(save) { + t.Fatalf("budget bytes = %d, want %d (deferred only)", bytes, record.Size(save)) + } + if got := len(drainUpload(in)); got != 0 { + t.Fatalf("upload got %d, want 0 before close", got) + } + + if err = in.Close(context.Background()); err != nil { + t.Fatalf("Close: %v", err) + } + if got := recordNums(drainUpload(in)); len(got) != 1 || got[0] != 1 { + t.Fatalf("upload after close = %v, want [1]", got) + } +} + +func TestCommitUploadOnlyWaitsForBudget(t *testing.T) { + b := budget.New(budget.Limits{Records: 1}) + if !b.TryAcquire(1, 0) { + t.Fatal("seed TryAcquire = false") + } + in := newIntake(t, Options{SkipStore: true, Budget: b, UploadQueue: 2}) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + if err := in.CommitUploadOnly(ctx, logRecord(1)); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("CommitUploadOnly err = %v, want context.DeadlineExceeded", err) + } + if records, _ := b.Usage(); records != 1 { + t.Fatalf("budget usage = %d, want 1", records) + } + if got := len(drainUpload(in)); got != 0 { + t.Fatalf("upload got %d, want 0", got) + } +} + +func TestCommitChargesActualBytes(t *testing.T) { + sink := newStoreSink(4) + defer sink.stop() + b := budget.New(budget.Limits{Bytes: 1 << 20}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 4}) + + rec := scalarRecord(1) + // 故意低估预留字节,Commit 必须按实际序列化大小补扣。 + tk, err := in.Reserve(context.Background(), Demand{Store: 1, Upload: 1, Bytes: 1}) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + if err := in.Commit(tk, []*recordv1.Record{rec}); err != nil { + t.Fatalf("Commit: %v", err) + } + _, bytes := b.Usage() + if want := record.Size(rec); bytes != want { + t.Fatalf("budget bytes = %d, want %d", bytes, want) + } +} + +func TestPendingEndSaveUploadsAtClose(t *testing.T) { + sink := newStoreSink(4) + defer sink.stop() + b := budget.New(budget.Limits{}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 4}) + + save := []*recordv1.Record{saveRecord(1, savev1.SavePolicy_SAVE_POLICY_END)} + mustCommit(t, in, save...) + if got := len(drainUpload(in)); got != 0 { + t.Fatalf("upload before close got %d, want 0", got) + } + if err := in.Close(context.Background()); err != nil { + t.Fatalf("Close: %v", err) + } + if got := len(drainUpload(in)); got != 1 { + t.Fatalf("upload after close got %d, want 1", got) + } + if records, _ := b.Usage(); records != 1 { + t.Fatalf("budget usage = %d, want 1", records) + } +} + +func TestCommitRoutes(t *testing.T) { + tests := []struct { + name string + records []*recordv1.Record + demand Demand + wantStore int + wantUpload int + flush bool // Close 后再统计上传通道 + }{ + { + name: "scalar both channels", + records: []*recordv1.Record{scalarRecord(1)}, + wantStore: 1, + wantUpload: 1, + }, + { + name: "start local only", + records: []*recordv1.Record{startRecord()}, + demand: Demand{Store: 1}, + wantStore: 1, + }, + { + name: "finish local only", + records: []*recordv1.Record{finishRecord()}, + demand: Demand{Store: 1}, + wantStore: 1, + }, + { + name: "log both channels", + records: []*recordv1.Record{logRecord(5)}, + wantStore: 1, + wantUpload: 1, + }, + { + name: "save now both channels", + records: []*recordv1.Record{saveRecord(6, savev1.SavePolicy_SAVE_POLICY_NOW)}, + wantStore: 1, + wantUpload: 1, + }, + { + name: "save end deferred until close", + records: []*recordv1.Record{saveRecord(7, savev1.SavePolicy_SAVE_POLICY_END)}, + wantStore: 1, + flush: true, + // 延迟上传仍占用上传预算,Close 后才进入通道。 + wantUpload: 1, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 8}) + + d := tc.demand + if len(tc.records) > 0 && (d.Store == 0 && d.Upload == 0 && d.Deferred == 0) { + d = demandFor(tc.records) + } + tk, err := in.Reserve(context.Background(), d) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + if err := in.Commit(tk, tc.records); err != nil { + t.Fatalf("Commit: %v", err) + } + + if got := waitStore(t, sink, tc.wantStore); len(got) != tc.wantStore { + t.Fatalf("store got %d, want %d", len(got), tc.wantStore) + } + if !tc.flush { + if got := len(drainUpload(in)); got != tc.wantUpload { + t.Fatalf("upload got %d, want %d", got, tc.wantUpload) + } + return + } + if got := len(drainUpload(in)); got != 0 { + t.Fatalf("upload before close got %d, want 0", got) + } + if err := in.Close(context.Background()); err != nil { + t.Fatalf("Close: %v", err) + } + if got := len(drainUpload(in)); got != tc.wantUpload { + t.Fatalf("upload after close got %d, want %d", got, tc.wantUpload) + } + }) + } +} + +func TestSkipStoreRoutesDataToUpload(t *testing.T) { + b := budget.New(budget.Limits{}) + in := newIntake(t, Options{SkipStore: true, Budget: b, UploadQueue: 4}) + recs := []*recordv1.Record{scalarRecord(1), scalarRecord(2)} + tk, err := in.Reserve(context.Background(), Demand{Upload: 2, Bytes: 32}) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + if err = in.Commit(tk, recs); err != nil { + t.Fatalf("Commit: %v", err) + } + if got := len(drainUpload(in)); got != 2 { + t.Fatalf("upload got %d, want 2", got) + } + if err := in.Close(context.Background()); err != nil { + t.Fatalf("Close: %v", err) + } +} + +func TestPauseAndResume(t *testing.T) { + sink := newStoreSink(16) + defer sink.stop() + b := budget.New(budget.Limits{Records: 2}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 8}) + + first := []*recordv1.Record{scalarRecord(1), scalarRecord(2)} + tk, err := in.Reserve(context.Background(), demandFor(first)) + if err != nil { + t.Fatalf("Reserve first: %v", err) + } + if err = in.Commit(tk, first); err != nil { + t.Fatalf("Commit first: %v", err) + } + + second := []*recordv1.Record{scalarRecord(3), scalarRecord(4)} + start := time.Now() + tk, err = in.Reserve(context.Background(), demandFor(second)) + if err != nil { + t.Fatalf("Reserve second: %v", err) + } + if elapsed := time.Since(start); elapsed > time.Second { + t.Fatalf("Reserve blocked for %v while a local copy exists", elapsed) + } + if err = in.Commit(tk, second); err != nil { + t.Fatalf("Commit second: %v", err) + } + + paused, seq := in.PauseState() + if !paused || seq != 3 { + t.Fatalf("PauseState = (%v, %d), want (true, 3)", paused, seq) + } + if got := waitStore(t, sink, 4); len(got) != 4 { + t.Fatalf("store got %d, want 4", len(got)) + } + if got := len(drainUpload(in)); got != 2 { + t.Fatalf("upload got %d, want 2 (paused batch must not upload)", got) + } + if records, _ := b.Usage(); records != 2 { + t.Fatalf("budget usage = %d, want 2", records) + } + + // 模拟上传确认后归还配额,再恢复投递。 + records, bytes := b.Usage() + b.Release(records, bytes) + in.ResumeUpload() + if paused, seq := in.PauseState(); paused || seq != 0 { + t.Fatalf("PauseState after resume = (%v, %d), want (false, 0)", paused, seq) + } + + third := []*recordv1.Record{scalarRecord(5)} + tk, err = in.Reserve(context.Background(), demandFor(third)) + if err != nil { + t.Fatalf("Reserve third: %v", err) + } + if err := in.Commit(tk, third); err != nil { + t.Fatalf("Commit third: %v", err) + } + if got := len(drainUpload(in)); got != 1 { + t.Fatalf("upload after resume got %d, want 1", got) + } + if records, _ := b.Usage(); records != 1 { + t.Fatalf("budget usage = %d, want 1", records) + } +} + +func TestPausedTicketAfterResumeStartsNewInterval(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{Records: 1}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 4}) + + first := []*recordv1.Record{scalarRecord(1)} + mustCommit(t, in, first...) + if got := len(drainUpload(in)); got != 1 { + t.Fatalf("upload got %d, want 1", got) + } + + second := []*recordv1.Record{scalarRecord(2)} + tk := mustReserve(t, in, second...) + // 回填完成并恢复投递,旧票据随后才提交:必须开启新的回填区间, + // 不能既不投递又不记录 pauseSeq。 + in.ResumeUpload() + if err := in.Commit(tk, second); err != nil { + t.Fatalf("Commit after resume: %v", err) + } + if got := len(drainUpload(in)); got != 0 { + t.Fatalf("upload got %d, want 0 (record must wait for backfill)", got) + } + if paused, seq := in.PauseState(); !paused || seq != 2 { + t.Fatalf("PauseState = (%v, %d), want (true, 2)", paused, seq) + } + if got := waitStore(t, sink, 2); len(got) != 2 { + t.Fatalf("store got %d, want 2", len(got)) + } + + // 模拟第二次回填完成后恢复,新记录应正常上传。 + records, bytes := b.Usage() + b.Release(records, bytes) + in.ResumeUpload() + third := []*recordv1.Record{scalarRecord(3)} + mustCommit(t, in, third...) + if got := len(drainUpload(in)); got != 1 { + t.Fatalf("upload after second resume got %d, want 1", got) + } + if paused, seq := in.PauseState(); paused || seq != 0 { + t.Fatalf("PauseState = (%v, %d), want (false, 0)", paused, seq) + } +} + +// TestAbandonedPausedTicketLeavesNoGlobalPause 覆盖暂停票据未成功提交的路径: +// 主动放弃、空提交、nil 记录与超预留都只能归还容量,不能遗留全局暂停状态。 +func TestAbandonedPausedTicketLeavesNoGlobalPause(t *testing.T) { + tests := []struct { + name string + abandon func(t *testing.T, in *Intake, tk Ticket) + }{ + { + name: "release", + abandon: func(_ *testing.T, in *Intake, tk Ticket) { in.Release(tk) }, + }, + { + name: "empty commit", + abandon: func(t *testing.T, in *Intake, tk Ticket) { + if err := in.Commit(tk, nil); err != nil { + t.Fatalf("Commit empty: %v", err) + } + }, + }, + { + name: "nil record", + abandon: func(t *testing.T, in *Intake, tk Ticket) { + if err := in.Commit(tk, []*recordv1.Record{nil}); err == nil { + t.Fatal("Commit nil record returned nil error") + } + }, + }, + { + name: "over reserved", + abandon: func(t *testing.T, in *Intake, tk Ticket) { + err := in.Commit(tk, []*recordv1.Record{scalarRecord(1), scalarRecord(2), scalarRecord(3)}) + if !errors.Is(err, ErrOverReserved) { + t.Fatalf("Commit err = %v, want ErrOverReserved", err) + } + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{Records: 1}) + if !b.TryAcquire(1, 0) { + t.Fatal("seed TryAcquire = false") + } + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 4}) + + tk, err := in.Reserve(context.Background(), Demand{Store: 2, Upload: 2, Bytes: 20}) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + if !tk.paused { + t.Fatal("ticket should be paused while the budget is full") + } + tc.abandon(t, in, tk) + if paused, seq := in.PauseState(); paused || seq != 0 { + t.Fatalf("PauseState = (%v, %d), want (false, 0)", paused, seq) + } + if got := len(sink.snapshot()); got != 0 { + t.Fatalf("store received %d records, want 0", got) + } + + // 预算空闲后下一次提交必须重新尝试获取配额并正常上传。 + b.Release(1, 0) + mustCommit(t, in, scalarRecord(1)) + if got := len(drainUpload(in)); got != 1 { + t.Fatalf("upload got %d, want 1", got) + } + if paused, seq := in.PauseState(); paused || seq != 0 { + t.Fatalf("PauseState after commit = (%v, %d), want (false, 0)", paused, seq) + } + }) + } +} + +func TestDegradedUsesBoundedBackpressure(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{Records: 1}) + if !b.TryAcquire(1, 0) { + t.Fatal("seed TryAcquire = false") + } + var degraded atomic.Bool + degraded.Store(true) + in := newIntake(t, Options{ + Store: sink.ch, + Budget: b, + Degraded: °raded, + UploadQueue: 4, + QueueWait: 30 * time.Millisecond, + }) + + start := time.Now() + _, err := in.Reserve(context.Background(), Demand{Store: 1, Upload: 1, Bytes: 10}) + elapsed := time.Since(start) + if !errors.Is(err, ErrOverloaded) { + t.Fatalf("Reserve err = %v, want ErrOverloaded", err) + } + if elapsed < 20*time.Millisecond { + t.Fatalf("Reserve returned in %v, want bounded wait", elapsed) + } + if records, _ := b.Usage(); records != 1 { + t.Fatalf("budget usage = %d, want 1 (failed reserve must not change state)", records) + } + if got := len(sink.snapshot()); got != 0 { + t.Fatalf("store received %d records, want 0", got) + } + if got := len(drainUpload(in)); got != 0 { + t.Fatalf("upload got %d records, want 0", got) + } +} + +func TestStoreQueueFullBackpressuresWithoutLoss(t *testing.T) { + storeCh := make(chan store.Request, 4) + for i := 0; i < cap(storeCh); i++ { + storeCh <- store.Request{StoreSeq: uint64(i + 1), Payload: []byte("x")} + } + var degraded atomic.Bool + b := budget.New(budget.Limits{}) + in := newIntake(t, Options{ + Store: storeCh, + Budget: b, + Degraded: °raded, + UploadQueue: 8, + QueueWait: 20 * time.Millisecond, + StorePollInterval: time.Millisecond, + }) + + recs := []*recordv1.Record{scalarRecord(1), scalarRecord(2)} + start := time.Now() + _, err := in.Reserve(context.Background(), demandFor(recs)) + if !errors.Is(err, ErrOverloaded) { + t.Fatalf("Reserve err = %v, want ErrOverloaded", err) + } + if elapsed := time.Since(start); elapsed < 15*time.Millisecond { + t.Fatalf("Reserve returned in %v, want bounded wait", elapsed) + } + if degraded.Load() { + t.Fatal("degraded = true after capacity pressure, want false (no loss/降级)") + } + if records, bytes := b.Usage(); records != 0 || bytes != 0 { + t.Fatalf("budget usage = (%d, %d), want (0, 0)", records, bytes) + } + if got := len(drainUpload(in)); got != 0 { + t.Fatalf("upload got %d records, want 0", got) + } + if got := len(storeCh); got != 4 { + t.Fatalf("store queue length = %d, want 4 (no writes)", got) + } +} + +func TestBatchLargerThanStoreQueueIsRejected(t *testing.T) { + storeCh := make(chan store.Request, 2) + var degraded atomic.Bool + in := newIntake(t, Options{ + Store: storeCh, + Budget: budget.New(budget.Limits{}), + Degraded: °raded, + UploadQueue: 4, + }) + _, err := in.Reserve(context.Background(), Demand{Store: 3, Upload: 3, Bytes: 30}) + if !errors.Is(err, ErrBatchTooLarge) { + t.Fatalf("Reserve err = %v, want ErrBatchTooLarge", err) + } + if degraded.Load() { + t.Fatal("degraded = true, want false") + } + if got := len(storeCh); got != 0 { + t.Fatalf("store queue length = %d, want 0", got) + } +} + +func TestPausedBatchFallsBackToUploadWhenStoreDegrades(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{Records: 1}) + var degraded atomic.Bool + in := newIntake(t, Options{Store: sink.ch, Budget: b, Degraded: °raded, UploadQueue: 4}) + + first := []*recordv1.Record{scalarRecord(1)} + mustCommit(t, in, first...) + if got := len(drainUpload(in)); got != 1 { + t.Fatalf("upload got %d, want 1", got) + } + + // 预算满 → 第二批按暂停处理;随后本地存储降级,Commit 必须改走上传。 + second := []*recordv1.Record{scalarRecord(2)} + tk := mustReserve(t, in, second...) + degraded.Store(true) + if err := in.Commit(tk, second); err != nil { + t.Fatalf("Commit second: %v", err) + } + if got := len(drainUpload(in)); got != 1 { + t.Fatalf("upload after degrade got %d, want 1 (paused batch must not be dropped)", got) + } + // 运行期降级只影响 local 判断,不得回写启动期配置。 + if in.skipStore { + t.Fatal("storeDegraded must not rewrite the skip_store setting") + } + if in.storeOut == nil { + t.Fatal("store channel must stay configured after degrade") + } + if got := waitStore(t, sink, 1); len(got) != 1 { + t.Fatalf("store got %d, want 1 (only the first record)", len(got)) + } + if records, _ := b.Usage(); records != 2 { + t.Fatalf("budget usage = %d, want 2 (temporary overrun is allowed to avoid loss)", records) + } +} + +func TestCommitUploadOnly(t *testing.T) { + sink := newStoreSink(4) + defer sink.stop() + b := budget.New(budget.Limits{}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 4}) + + rec := logRecord(9) + if err := in.CommitUploadOnly(context.Background(), rec); err != nil { + t.Fatalf("CommitUploadOnly: %v", err) + } + if got := len(drainUpload(in)); got != 1 { + t.Fatalf("upload got %d, want 1", got) + } + if got := len(sink.snapshot()); got != 0 { + t.Fatalf("store received %d records, want 0", got) + } + if records, _ := b.Usage(); records != 1 { + t.Fatalf("budget usage = %d, want 1", records) + } +} + +func TestCommitUploadOnlyAfterClose(t *testing.T) { + b := budget.New(budget.Limits{}) + in := newIntake(t, Options{SkipStore: true, Budget: b, UploadQueue: 2}) + if err := in.Close(context.Background()); err != nil { + t.Fatalf("Close: %v", err) + } + err := in.CommitUploadOnly(context.Background(), logRecord(1)) + if !errors.Is(err, ErrClosed) { + t.Fatalf("CommitUploadOnly err = %v, want ErrClosed", err) + } + if records, _ := b.Usage(); records != 0 { + t.Fatalf("budget usage = %d, want 0 after rejected commit", records) + } +} + +func TestCloseIsIdempotentAndRejectsCommits(t *testing.T) { + sink := newStoreSink(4) + defer sink.stop() + b := budget.New(budget.Limits{Records: 4}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 4}) + + tk, err := in.Reserve(context.Background(), Demand{Store: 1, Upload: 1, Bytes: 10}) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + if err = in.Close(context.Background()); err != nil { + t.Fatalf("Close: %v", err) + } + if err = in.Close(context.Background()); err != nil { + t.Fatalf("second Close: %v", err) + } + if err = in.Commit(tk, []*recordv1.Record{scalarRecord(1)}); !errors.Is(err, ErrClosed) { + t.Fatalf("Commit err = %v, want ErrClosed", err) + } + if records, bytes := b.Usage(); records != 0 || bytes != 0 { + t.Fatalf("budget usage = (%d, %d), want (0, 0)", records, bytes) + } + if _, err = in.Reserve(context.Background(), Demand{Store: 1, Upload: 1}); !errors.Is(err, ErrClosed) { + t.Fatalf("Reserve after close err = %v, want ErrClosed", err) + } +} + +// TestCloseRestoresPendingWhenRecoverySendingFails 覆盖 Close 在补投回收记录时 +// 被取消:已暂存的 END 保存不能随快照一起丢失,重试 Close 必须补发。 +func TestCloseRestoresPendingWhenRecoverySendingFails(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{Records: 1}) + var degraded atomic.Bool + in := newIntake(t, Options{ + Store: sink.ch, + Budget: b, + Degraded: °raded, + UploadQueue: 1, + Flushed: func() uint64 { return 0 }, + }) + + mustCommit(t, in, scalarRecord(1)) // 占满预算并填满上传通道 + mustCommit(t, in, scalarRecord(2)) // 暂停,尚未确认落盘 + b.Release(1, record.Size(scalarRecord(1))) + mustCommit(t, in, saveRecord(3, savev1.SavePolicy_SAVE_POLICY_END)) + + degraded.Store(true) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + if err := in.Close(ctx); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("Close err = %v, want context.DeadlineExceeded", err) + } + if got := recordNums(drainUpload(in)); len(got) != 1 || got[0] != 1 { + t.Fatalf("upload = %v, want [1]", got) + } + + var got []int64 + consumed := make(chan struct{}) + go func() { + defer close(consumed) + for rec := range in.UploadOut() { + got = append(got, rec.GetNum()) + } + }() + if err := in.Close(context.Background()); err != nil { + t.Fatalf("second Close: %v", err) + } + <-consumed + if len(got) != 2 || got[0] != 2 || got[1] != 3 { + t.Fatalf("upload after retry = %v, want [2 3]", got) + } +} + +func TestOverReservedCommitHasNoSideEffects(t *testing.T) { + sink := newStoreSink(4) + defer sink.stop() + b := budget.New(budget.Limits{}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 4}) + + tk, err := in.Reserve(context.Background(), Demand{Store: 1, Upload: 1, Bytes: 10}) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + err = in.Commit(tk, []*recordv1.Record{scalarRecord(1), scalarRecord(2)}) + if !errors.Is(err, ErrOverReserved) { + t.Fatalf("Commit err = %v, want ErrOverReserved", err) + } + if got := len(sink.snapshot()); got != 0 { + t.Fatalf("store received %d records, want 0", got) + } + if got := len(drainUpload(in)); got != 0 { + t.Fatalf("upload got %d records, want 0", got) + } + if records, bytes := b.Usage(); records != 0 || bytes != 0 { + t.Fatalf("budget usage = (%d, %d), want (0, 0)", records, bytes) + } +} + +func TestReleaseUnusedReservation(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{Records: 4}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 4}) + + recs := []*recordv1.Record{scalarRecord(1), scalarRecord(2)} + tk, err := in.Reserve(context.Background(), Demand{Store: 3, Upload: 3, Bytes: 30}) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + if err = in.Commit(tk, recs); err != nil { + t.Fatalf("Commit: %v", err) + } + if records, _ := b.Usage(); records != 2 { + t.Fatalf("budget usage = %d, want 2", records) + } + // 预留的 1 条本地空间也必须归还,否则后续 Reserve 会误判队列已满。 + tk2, err := in.Reserve(context.Background(), Demand{Store: 2, Upload: 2, Bytes: 20}) + if err != nil { + t.Fatalf("second Reserve: %v", err) + } + in.Release(tk2) + if got := waitStore(t, sink, 2); len(got) != 2 { + t.Fatalf("store got %d, want 2", len(got)) + } +} + +func TestConcurrentCommits(t *testing.T) { + const workers = 8 + sink := newStoreSink(workers) + defer sink.stop() + b := budget.New(budget.Limits{}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: workers}) + + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + recs := []*recordv1.Record{scalarRecord(int64(i + 1))} + tk, err := in.Reserve(context.Background(), demandFor(recs)) + if err != nil { + t.Errorf("Reserve: %v", err) + return + } + if err = in.Commit(tk, recs); err != nil { + t.Errorf("Commit: %v", err) + } + }(i) + } + wg.Wait() + if err := in.Close(context.Background()); err != nil { + t.Fatalf("Close: %v", err) + } + + got := waitStore(t, sink, workers) + if len(got) != workers { + t.Fatalf("store got %d, want %d", len(got), workers) + } + seen := make(map[uint64]bool, workers) + for _, req := range got { + if req.StoreSeq == 0 || req.StoreSeq > workers { + t.Fatalf("StoreSeq = %d, want 1..%d", req.StoreSeq, workers) + } + if seen[req.StoreSeq] { + t.Fatalf("duplicate StoreSeq %d", req.StoreSeq) + } + seen[req.StoreSeq] = true + } + if got := len(drainUpload(in)); got != workers { + t.Fatalf("upload got %d, want %d", got, workers) + } +} + +// TestPausedRecordsRecoverAfterStoreFailure 注入真实写盘失败:Commit 成功返回后 +// 关闭底层 Writer,AsyncWriter 写入/刷盘失败并置 storeDegraded;暂停期间未确认 +// 落盘的记录必须由 DrainRecovery 重新投入上传,不能丢。 +func TestPausedRecordsRecoverAfterStoreFailure(t *testing.T) { + w, err := store.Create(filepath.Join(t.TempDir(), "run-recover.swanlab")) + if err != nil { + t.Fatalf("store.Create: %v", err) + } + var degraded atomic.Bool + aw, err := store.NewAsyncWriter(w, store.AsyncOptions{ + Queue: 8, + FlushInterval: time.Hour, + FlushBytes: 1 << 30, + OnError: func(error) { degraded.Store(true) }, + }) + if err != nil { + t.Fatalf("NewAsyncWriter: %v", err) + } + defer func() { _ = aw.Close() }() + + b := budget.New(budget.Limits{Records: 1}) + in := newIntake(t, Options{ + Store: aw.In(), + Budget: b, + Degraded: °raded, + UploadQueue: 8, + Flushed: func() uint64 { + seq, _ := aw.Flushed() + return seq + }, + }) + + first := []*recordv1.Record{scalarRecord(1)} + tk, err := in.Reserve(context.Background(), demandFor(first)) + if err != nil { + t.Fatalf("Reserve first: %v", err) + } + if err = in.Commit(tk, first); err != nil { + t.Fatalf("Commit first: %v", err) + } + if got := len(drainUpload(in)); got != 1 { + t.Fatalf("upload got %d, want 1", got) + } + + // 预算耗尽 → 第二批只落本地并保留未确认引用。 + second := []*recordv1.Record{scalarRecord(2)} + tk, err = in.Reserve(context.Background(), demandFor(second)) + if err != nil { + t.Fatalf("Reserve second: %v", err) + } + if err = in.Commit(tk, second); err != nil { + t.Fatalf("Commit second: %v", err) + } + if paused, seq := in.PauseState(); !paused || seq != 2 { + t.Fatalf("PauseState = (%v, %d), want (true, 2)", paused, seq) + } + + // 真实故障:关闭底层文件,强制 AsyncWriter 在后续写入/刷盘时报错。 + if err = w.Close(); err != nil { + t.Fatalf("close writer: %v", err) + } + _ = aw.FlushForTest() + deadline := time.After(time.Second) + for !degraded.Load() { + select { + case <-deadline: + t.Fatal("storeDegraded was not set after write failure") + default: + time.Sleep(time.Millisecond) + } + } + + remaining, err := in.DrainRecovery(context.Background()) + if err != nil { + t.Fatalf("DrainRecovery: %v", err) + } + if remaining != 0 { + t.Fatalf("recovery remaining = %d, want 0", remaining) + } + got := drainUpload(in) + if len(got) != 1 || got[0].GetNum() != 2 { + t.Fatalf("recovered upload = %v, want record num 2", recordNums(got)) + } +} + +// TestPausedThenDegradedUsesBudgetBackpressure 覆盖“先暂停、后降级”: +// 降级后不得再走 paused 快路径,必须回到有界预算反压。 +func TestPausedThenDegradedUsesBudgetBackpressure(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{Records: 1}) + if !b.TryAcquire(1, 0) { + t.Fatal("seed TryAcquire = false") + } + var degraded atomic.Bool + in := newIntake(t, Options{ + Store: sink.ch, + Budget: b, + Degraded: °raded, + UploadQueue: 4, + QueueWait: 30 * time.Millisecond, + }) + + first := []*recordv1.Record{scalarRecord(1)} + mustCommit(t, in, first...) + if paused, seq := in.PauseState(); !paused || seq != 1 { + t.Fatalf("PauseState = (%v, %d), want (true, 1)", paused, seq) + } + + degraded.Store(true) + start := time.Now() + _, err := in.Reserve(context.Background(), demandFor([]*recordv1.Record{scalarRecord(2)})) + elapsed := time.Since(start) + if !errors.Is(err, ErrOverloaded) { + t.Fatalf("Reserve err = %v, want ErrOverloaded", err) + } + if elapsed < 20*time.Millisecond { + t.Fatalf("Reserve returned in %v, want bounded wait", elapsed) + } + if got := len(sink.snapshot()); got != 1 { + t.Fatalf("store got %d, want 1 (rejected batch must not write)", got) + } + if records, _ := b.Usage(); records != 1 { + t.Fatalf("budget usage = %d, want 1 (failed reserve must not change state)", records) + } + + // 释放预算后,降级路径应能正常占用配额并只走上传。 + records, bytes := b.Usage() + b.Release(records, bytes) + mustCommit(t, in, scalarRecord(3)) + if got := len(drainUpload(in)); got != 1 { + t.Fatalf("upload got %d, want 1", got) + } +} + +// TestUploadQueueFullBackpressuresReserve 验证上传通道满时在 Reserve 阶段 +// 有界反压,而不是在 Commit 持锁无界阻塞。 +func TestUploadQueueFullBackpressuresReserve(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{}) + in := newIntake(t, Options{ + Store: sink.ch, + Budget: b, + UploadQueue: 1, + QueueWait: 25 * time.Millisecond, + }) + + first := []*recordv1.Record{scalarRecord(1)} + tk, err := in.Reserve(context.Background(), demandFor(first)) + if err != nil { + t.Fatalf("Reserve first: %v", err) + } + if err = in.Commit(tk, first); err != nil { + t.Fatalf("Commit first: %v", err) + } + + start := time.Now() + _, err = in.Reserve(context.Background(), demandFor([]*recordv1.Record{scalarRecord(2)})) + elapsed := time.Since(start) + if !errors.Is(err, ErrOverloaded) { + t.Fatalf("Reserve err = %v, want ErrOverloaded", err) + } + if elapsed < 20*time.Millisecond || elapsed > time.Second { + t.Fatalf("Reserve blocked for %v, want bounded wait", elapsed) + } + + done := make(chan error, 1) + go func() { done <- in.Close(context.Background()) }() + select { + case err := <-done: + if err != nil { + t.Fatalf("Close: %v", err) + } + case <-time.After(time.Second): + t.Fatal("Close blocked on upload queue") + } +} + +// TestPausedBatchUsesRecoveryWhenUploadQueueFull 覆盖暂停票据在降级后 +// 上传通道也满的情况:不得持锁阻塞,记录暂存回收队列,等通道腾空后补传。 +func TestPausedBatchUsesRecoveryWhenUploadQueueFull(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{Records: 1}) + var degraded atomic.Bool + in := newIntake(t, Options{ + Store: sink.ch, + Budget: b, + Degraded: °raded, + UploadQueue: 1, + }) + + first := []*recordv1.Record{scalarRecord(1)} + mustCommit(t, in, first...) + + second := []*recordv1.Record{scalarRecord(2)} + tk := mustReserve(t, in, second...) + degraded.Store(true) + if err := in.Commit(tk, second); err != nil { + t.Fatalf("Commit second: %v", err) + } + if n, err := in.DrainRecovery(context.Background()); n != 1 || err != nil { + t.Fatalf("DrainRecovery = (%d, %v), want (1, nil) while channel is full", n, err) + } + + if got := len(drainUpload(in)); got != 1 { + t.Fatalf("upload got %d, want 1 (only the first record)", got) + } + if n, err := in.DrainRecovery(context.Background()); n != 0 || err != nil { + t.Fatalf("DrainRecovery = (%d, %v), want (0, nil)", n, err) + } + got := drainUpload(in) + if len(got) != 1 || got[0].GetNum() != 2 { + t.Fatalf("recovered upload = %v, want record num 2", recordNums(got)) + } + if records, _ := b.Usage(); records != 2 { + t.Fatalf("budget usage = %d, want 2", records) + } + if got := waitStore(t, sink, 1); len(got) != 1 { + t.Fatalf("store got %d, want 1 (only the first record)", len(got)) + } +} + +// TestRecoveryKeepsUploadOrderAfterDegrade 覆盖“暂停记录 + 后续正常预留批次”: +// 降级补传必须先把更早的暂停记录送出,后续批次不能反超。 +func TestRecoveryKeepsUploadOrderAfterDegrade(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{Records: 2}) + var degraded atomic.Bool + in := newIntake(t, Options{ + Store: sink.ch, + Budget: b, + Degraded: °raded, + UploadQueue: 8, + // 暂停记录尚未确认落盘,降级后必须重新投入上传。 + Flushed: func() uint64 { return 0 }, + }) + + mustCommit(t, in, scalarRecord(1), scalarRecord(2)) + if got := recordNums(drainUpload(in)); len(got) != 2 || got[0] != 1 || got[1] != 2 { + t.Fatalf("upload = %v, want [1 2]", got) + } + + // 预算已满 → 记录 3 暂停,只落本地并保留未确认引用。 + mustCommit(t, in, scalarRecord(3)) + if paused, seq := in.PauseState(); !paused || seq != 3 { + t.Fatalf("PauseState = (%v, %d), want (true, 3)", paused, seq) + } + + // 模拟上传确认一条并归还配额;本地随后降级。 + b.Release(1, record.Size(scalarRecord(1))) + degraded.Store(true) + mustCommit(t, in, scalarRecord(4)) + + if got := recordNums(drainUpload(in)); len(got) != 2 || got[0] != 3 || got[1] != 4 { + t.Fatalf("upload after degrade = %v, want [3 4]", got) + } +} + +// TestPauseKeepsDeferredQuota 覆盖“暂停前预留、暂停后提交”的票据: +// 暂停只归还立即记录的配额,延迟记录继续持有到 Close。 +func TestPauseKeepsDeferredQuota(t *testing.T) { + sink := newStoreSink(8) + defer sink.stop() + b := budget.New(budget.Limits{Records: 2}) + in := newIntake(t, Options{Store: sink.ch, Budget: b, UploadQueue: 4}) + + save := saveRecord(2, savev1.SavePolicy_SAVE_POLICY_END) + batch := []*recordv1.Record{scalarRecord(1), save} + tk, err := in.Reserve(context.Background(), demandFor(batch)) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + + // 票据占满预算 → 另一批暂停,随后旧票据才提交。 + mustCommit(t, in, scalarRecord(3)) + if err = in.Commit(tk, batch); err != nil { + t.Fatalf("Commit deferred batch: %v", err) + } + if records, _ := b.Usage(); records != 1 { + t.Fatalf("budget usage = %d, want 1 (only the deferred save stays charged)", records) + } + if got := len(drainUpload(in)); got != 0 { + t.Fatalf("upload got %d, want 0 before close", got) + } + + if err = in.Close(context.Background()); err != nil { + t.Fatalf("Close: %v", err) + } + got := drainUpload(in) + if len(got) != 1 || got[0].GetNum() != 2 { + t.Fatalf("upload after close = %v, want record num 2", recordNums(got)) + } + if records, _ := b.Usage(); records != 1 { + t.Fatalf("budget usage = %d, want 1", records) + } +} + +func recordNums(recs []*recordv1.Record) []int64 { + nums := make([]int64, len(recs)) + for i, rec := range recs { + nums[i] = rec.GetNum() + } + return nums +} diff --git a/core/internal/buffer/recovery.go b/core/internal/buffer/recovery.go new file mode 100644 index 000000000..22aae3f3e --- /dev/null +++ b/core/internal/buffer/recovery.go @@ -0,0 +1,180 @@ +// 本文件承载暂停投递后的未确认记录、回填区间与本地失效后的回收补传。 +package buffer + +import ( + "context" + + "github.com/swanhubx/swanlab/core/internal/pkg/record" + recordv1 "github.com/swanhubx/swanlab/core/proto/swanlab/record/v1" +) + +// unconfirmedRec 是已进入本地写入队列、但尚未确认刷盘的暂停记录。 +// 本地写失败时依赖它把记录重新投入上传,避免既无本地副本又不在上传管线。 +type unconfirmedRec struct { + seq uint64 + rec *recordv1.Record +} + +// PauseState 返回暂停状态与首个未投递记录的 storeSeq(0 表示未暂停)。 +func (t *Intake) PauseState() (paused bool, storeSeq uint64) { + t.mu.Lock() + defer t.mu.Unlock() + return t.pauseSeq != 0, t.pauseSeq +} + +// ResumeUpload 恢复上传投递。调用方必须先补传 pauseSeq 之后的本地记录; +// 恢复后仍可能因旧票据或预算再次进入暂停,须重新观察 PauseState。 +func (t *Intake) ResumeUpload() { + t.mu.Lock() + t.pauseSeq = 0 + t.mu.Unlock() +} + +// DrainRecovery 尝试把回收队列投入上传;返回仍未投递的记录数。 +// 本地已降级时会把未确认落盘的暂停记录一并转入回收队列,可重复调用。 +func (t *Intake) DrainRecovery(ctx context.Context) (int, error) { + t.mu.Lock() + defer t.mu.Unlock() + if t.degraded.Load() { + t.collectUnconfirmedLocked() + } + if len(t.recovery) == 0 { + return 0, nil + } + if t.closed { + return len(t.recovery), ErrClosed + } + if err := ctx.Err(); err != nil { + return len(t.recovery), err + } + t.drainRecoveryLocked() + return len(t.recovery), nil +} + +// drainRecoveryLocked 按 FIFO 尝试补发回收记录;槽位不足时停在最早未投递的 +// 记录上,返回是否已全部送出。调用方需持有 mu。 +func (t *Intake) drainRecoveryLocked() bool { + for len(t.recovery) > 0 { + if !t.sendRecoveryLocked(t.recovery[0]) { + return false + } + t.recovery[0] = nil + t.recovery = t.recovery[1:] + } + t.recovery = nil + return true +} + +// pauseLocked 记录暂停起点、保留未确认记录、暂存 END 保存并归还本批上传槽位。 +// 延迟记录从 Reserve 起就占用配额,这里不释放;若 Resume 之后仍有旧票据提交, +// 这里会开启新的回填区间。返回是否新发布了暂停状态。需持有 mu。 +func (t *Intake) pauseLocked(records []*recordv1.Record, immediate, deferred []int, tk Ticket, seqOf []uint64) bool { + published := false + if t.pauseSeq == 0 && len(immediate) > 0 { + // 暂停状态只在暂停记录真正提交时发布,票据被放弃不会遗留暂停。 + t.pauseSeq = seqOf[immediate[0]] + published = true + } + for _, i := range immediate { + t.unconfirmed = append(t.unconfirmed, unconfirmedRec{seq: seqOf[i], rec: records[i]}) + } + for _, i := range deferred { + t.pending = append(t.pending, records[i]) + } + if tk.upload > 0 { + // 只归还立即记录的配额;延迟记录继续持有到 Close。 + t.budget.Adjust(int64(len(deferred))-int64(tk.upload), deferredBytes(records, deferred)-tk.bytes) + } + t.releaseUploadLocked(tk.slots) + return published +} + +// recoveryUploadLocked 处理提交时没有可靠本地副本的批次:并入回收队列, +// 排在已有记录之后按序补传,绝不越过更早的记录。槽位不足时留在队列中, +// 由后续 DrainRecovery 或 Close 继续。需持有 mu。 +func (t *Intake) recoveryUploadLocked(records []*recordv1.Record, immediate, deferred []int, tk Ticket) { + if tk.upload > 0 { + // 延迟记录继续持有配额;立即记录随回收队列在发送时补扣。 + t.budget.Adjust(int64(len(deferred))-int64(tk.upload), deferredBytes(records, deferred)-tk.bytes) + } + for _, i := range immediate { + t.recovery = append(t.recovery, records[i]) + } + for _, i := range deferred { + t.pending = append(t.pending, records[i]) + } + t.releaseUploadLocked(tk.slots) + t.drainRecoveryLocked() +} + +// deferredBytes 统计延迟记录的序列化字节数。 +func deferredBytes(records []*recordv1.Record, deferred []int) int64 { + var total int64 + for _, i := range deferred { + total += record.Size(records[i]) + } + return total +} + +// sendRecoveryLocked 为一条回收记录占用槽位、补扣配额并投入上传;需持有 mu。 +// 本地副本已失效,因此不再等待预算,超限部分由回填积压量约束。 +func (t *Intake) sendRecoveryLocked(rec *recordv1.Record) bool { + if !t.tryReserveUploadLocked(1) { + return false + } + t.budget.Adjust(1, record.Size(rec)) + t.upload <- rec + t.releaseUploadLocked(1) + return true +} + +// trimUnconfirmedLocked 丢弃已确认刷盘的前缀;调用方需持有 mu。 +func (t *Intake) trimUnconfirmedLocked() { + if len(t.unconfirmed) == 0 { + return + } + flushed := t.flushedSeq() + i := 0 + for i < len(t.unconfirmed) && t.unconfirmed[i].seq <= flushed { + t.unconfirmed[i] = unconfirmedRec{} + i++ + } + if i > 0 { + t.unconfirmed = t.unconfirmed[i:] + } + if len(t.unconfirmed) == 0 { + t.unconfirmed = nil + } +} + +// collectUnconfirmedLocked 把未确认落盘的暂停记录转入回收队列;调用方需持有 mu。 +func (t *Intake) collectUnconfirmedLocked() { + if len(t.unconfirmed) == 0 { + return + } + flushed := t.flushedSeq() + for _, u := range t.unconfirmed { + if u.seq > flushed { + t.recovery = append(t.recovery, u.rec) + } + } + t.unconfirmed = nil +} + +func (t *Intake) flushedSeq() uint64 { + if t.flushed == nil { + return 0 + } + return t.flushed() +} + +// restoreRecovery 把未发出的回收记录放回队头,保持原有顺序。 +func (t *Intake) restoreRecovery(items []*recordv1.Record) { + if len(items) == 0 { + return + } + t.mu.Lock() + rest := append([]*recordv1.Record(nil), items...) + t.recovery = append(rest, t.recovery...) + t.mu.Unlock() +} diff --git a/core/internal/config/settings.go b/core/internal/config/settings.go index ad8547d85..b4ccf83c7 100644 --- a/core/internal/config/settings.go +++ b/core/internal/config/settings.go @@ -73,6 +73,7 @@ type Settings struct { SaveBatch int32 // SkipStore 跳过本地 Record 落盘(仅 online 模式合法,由上游校验)。 + // 启动期配置,运行期不得修改;磁盘写满等故障属于会话运行期状态。 SkipStore bool // APIKey 云端认证凭据;日志输出必须脱敏。