fix(submit): 分块上传的预算检查复用一分钟内读过的 blob 大小，存储读不到时回 503 让面板自动重试

internal/submit/sizes_test.go

@@ -0,0 +1,175 @@
+package submit
+
+import (
+	"context"
+	"errors"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+)
+
+// countingBlobs counts the Size reads per id.
+type countingBlobs struct {
+	*fakeBlobs
+	sizes map[string]int
+}
+
+func (c *countingBlobs) Size(ctx context.Context, id string) (int64, bool, error) {
+	c.sizes[id]++
+	return c.fakeBlobs.Size(ctx, id)
+}
+
+// putThenFail stores the bytes and still fails, like an object-store upload that
+// completed but whose answer was lost.
+type putThenFail struct{ *fakeBlobs }
+
+func (p putThenFail) Put(ctx context.Context, id string, r io.Reader) (int64, error) {
+	if _, err := p.fakeBlobs.Put(ctx, id, r); err != nil {
+		return 0, err
+	}
+	return 0, errors.New("put: connection reset")
+}
+
+// The budget check ahead of each part reads every other blob's size once per
+// blobSizeTTL; the one on completion reads them all from the store.
+func TestPartBudgetRemembersBlobSizes(t *testing.T) {
+	m, _, fb, a := newChunkedManager(t)
+	ctx := context.Background()
+	clock := testNow
+	m.Now = func() time.Time { return clock }
+	b, err := m.Create(ctx, CreateRequest{DisplayName: "B", SubmittedBy: "user-2"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	fb.stored[b.ID] = []byte("xyz")
+	cb := &countingBlobs{fakeBlobs: fb, sizes: map[string]int{}}
+	m.Blobs = cb
+
+	sendPart(t, m, a, 0, "\x1f\x8b\x08\x00")
+	sendPart(t, m, a, 4, "abcd")
+	if n := cb.sizes[b.ID]; n != 1 {
+		t.Fatalf("two parts read the other blob's size %d times, want once", n)
+	}
+	clock = clock.Add(blobSizeTTL)
+	sendPart(t, m, a, 8, "ef")
+	if n := cb.sizes[b.ID]; n != 2 {
+		t.Fatalf("a part once the TTL ran out: %d reads in all, want 2", n)
+	}
+	if _, err := m.CompleteUpload(ctx, a, "user-1"); err != nil {
+		t.Fatal(err)
+	}
+	if n := cb.sizes[b.ID]; n != 3 {
+		t.Fatalf("completion within the TTL: %d reads in all, want 3 (it reads the store)", n)
+	}
+}
+
+// What this Manager stores or deletes counts in the next part's budget at once,
+// without waiting out the TTL.
+func TestPartBudgetSeesThisManagersOwnWritesAtOnce(t *testing.T) {
+	m, _, _, a := newChunkedManager(t)
+	ctx := context.Background()
+	m.MaxStoredBytesPerUser = 10
+	m.PartMaxBytes = 16
+	b, err := m.Create(ctx, CreateRequest{DisplayName: "B", SubmittedBy: "user-1"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	// The first part remembers B as holding nothing.
+	sendPart(t, m, a, 0, "\x1f\x8b")
+	if _, err := m.UploadContext(ctx, b.ID, "user-1", strings.NewReader("\x1f\x8b\x08\x00ab")); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := m.UploadPart(ctx, a, "user-1", 2, strings.NewReader("cde")); !errors.Is(err, ErrQuotaExceeded) {
+		t.Fatalf("5 bytes staged beside B's 6 under a 10-byte budget = %v, want ErrQuotaExceeded", err)
+	}
+	if _, err := m.Withdraw(ctx, b.ID, "user-1"); err != nil {
+		t.Fatal(err)
+	}
+	// Its row is gone, so nothing counts it; nor is its size kept, or the
+	// remembered sizes would grow with every submission ever deleted.
+	if _, kept := m.sizes[b.ID]; kept {
+		t.Fatal("a withdrawn submission's blob size is still remembered")
+	}
+	if got := sendPart(t, m, a, 2, "cdefgh"); got != 8 {
+		t.Fatalf("staged %d after B was withdrawn, want 8", got)
+	}
+}
+
+func TestPartBudgetSeesAReapedBlobAtOnce(t *testing.T) {
+	m, st, _, a := newChunkedManager(t)
+	ctx := context.Background()
+	m.MaxStoredBytesPerUser = 10
+	m.PartMaxBytes = 16
+	b, err := m.Create(ctx, CreateRequest{DisplayName: "B", SubmittedBy: "user-1"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := m.UploadContext(ctx, b.ID, "user-1", strings.NewReader("\x1f\x8b\x08\x00ab")); err != nil {
+		t.Fatal(err)
+	}
+	reviewed := testNow.Add(-48 * time.Hour)
+	st.subs[b.ID].Status = StatusRejected
+	st.subs[b.ID].ReviewedAt = &reviewed
+	if n, err := m.ReapRejected(ctx, 24*time.Hour); n != 1 || err != nil {
+		t.Fatalf("ReapRejected = %d, %v; want 1, nil", n, err)
+	}
+	if got := sendPart(t, m, a, 0, "\x1f\x8b\x08\x00abcd"); got != 8 {
+		t.Fatalf("staged %d after B's blob was reaped, want 8", got)
+	}
+}
+
+// A Put that failed may still have replaced the blob, so its size is read again.
+func TestPartBudgetRereadsABlobAfterAFailedPut(t *testing.T) {
+	m, _, fb, a := newChunkedManager(t)
+	ctx := context.Background()
+	m.MaxStoredBytesPerUser = 10
+	m.PartMaxBytes = 16
+	b, err := m.Create(ctx, CreateRequest{DisplayName: "B", SubmittedBy: "user-1"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := m.UploadContext(ctx, b.ID, "user-1", strings.NewReader("\x1f\x8b\x08\x00ab")); err != nil {
+		t.Fatal(err)
+	}
+	m.Blobs = putThenFail{fb}
+	if _, err := m.UploadContext(ctx, b.ID, "user-1", strings.NewReader("\x1f\x8b")); err == nil {
+		t.Fatal("setup: the failing Put succeeded")
+	}
+	m.Blobs = fb
+	if got := sendPart(t, m, a, 0, "\x1f\x8b\x08\x00abcd"); got != 8 {
+		t.Fatalf("staged %d beside B's 2 stored bytes, want 8", got)
+	}
+}
+
+// A size that cannot be read fails the check as ErrStoreUnavailable, which the
+// API answers 503 for the client to send again, never as a bare error (500).
+func TestBudgetReadFailuresAreRetryable(t *testing.T) {
+	m, _, fb, a := newChunkedManager(t)
+	ctx := context.Background()
+	if _, err := m.Create(ctx, CreateRequest{DisplayName: "B", SubmittedBy: "user-2"}); err != nil {
+		t.Fatal(err)
+	}
+	fb.sizeErr = errors.New("dial tcp: i/o timeout")
+	if _, err := m.UploadPart(ctx, a, "user-1", 0, strings.NewReader("\x1f\x8b\x08\x00")); !errors.Is(err, ErrStoreUnavailable) {
+		t.Fatalf("a part with the blob store down = %v, want ErrStoreUnavailable", err)
+	}
+	if _, err := m.UploadContext(ctx, a, "user-1", strings.NewReader("\x1f\x8b\x08\x00")); !errors.Is(err, ErrStoreUnavailable) {
+		t.Fatalf("an upload with the blob store down = %v, want ErrStoreUnavailable", err)
+	}
+	if _, ok := fb.stored[a]; ok {
+		t.Fatal("an upload whose budget could not be checked was stored")
+	}
+
+	fb.sizeErr = nil
+	notDir := filepath.Join(t.TempDir(), "parts")
+	if err := os.WriteFile(notDir, nil, 0o600); err != nil {
+		t.Fatal(err)
+	}
+	m.Parts = &PartStore{Dir: notDir}
+	if _, err := m.UploadContext(ctx, a, "user-1", strings.NewReader("\x1f\x8b\x08\x00")); !errors.Is(err, ErrStoreUnavailable) {
+		t.Fatalf("an upload with the staging directory unreadable = %v, want ErrStoreUnavailable", err)
+	}
+}