1// SPDX-License-Identifier: AGPL-3.0-or-later
2
3package usecase
4
5import (
6 "bytes"
7 "context"
8 "errors"
9 "io"
10 "strings"
11 "sync"
12 "testing"
13 "time"
14
15 "github.com/stretchr/testify/mock"
16 "github.com/stretchr/testify/require"
17
18 "bnema.dev/ermite/internal/domain"
19 "bnema.dev/ermite/internal/ports"
20 "bnema.dev/ermite/internal/ports/mocks"
21)
22
23// gateReader signals its first Read on started, then blocks until release is
24// closed or ctx ends; it then returns data (or err when set).
25type gateReader struct {
26 ctx context.Context
27 started chan struct{}
28 release chan struct{}
29 data []byte
30 err error
31 once sync.Once
32 reads int
33 r *bytes.Reader
34}
35
36func (g *gateReader) Read(p []byte) (int, error) {
37 g.reads++
38 g.once.Do(func() {
39 if g.started != nil {
40 close(g.started)
41 }
42 })
43 if g.release != nil {
44 select {
45 case <-g.release:
46 case <-g.ctx.Done():
47 return 0, g.ctx.Err()
48 }
49 }
50 if g.err != nil {
51 return 0, g.err
52 }
53 if g.r == nil {
54 g.r = bytes.NewReader(g.data)
55 }
56 return g.r.Read(p)
57}
58
59type slotFixture struct {
60 c *CI
61 secret string
62 job domain.CIJob
63 puts *int
64}
65
66func newSlotFixture(t *testing.T, slots int) *slotFixture {
67 t.Helper()
68 repo, err := domain.ParseRepoName("proj")
69 require.NoError(t, err)
70 secret := domain.JobTokenPrefix + strings.Repeat("a", domain.CITokenRandomLen)
71 hash := domain.HashTokenSecret(secret)
72 j := domain.CIJob{ID: "1-0-0", RunID: 1, RunAttempt: 1, Repo: repo, TokenHash: hash, Phase: domain.CIPhaseRunning}
73 a := domain.CIArtifact{ID: 7, Repo: repo, RunID: 1, Attempt: 1, JobID: j.ID, Assignment: hash, State: domain.CIArtifactPending}
74
75 store := mocks.NewCIStore(t)
76 store.EXPECT().FindJobByToken(mock.Anything, hash, mock.Anything).Return(j, nil).Maybe()
77 store.EXPECT().GetJob(mock.Anything, j.RunID, j.RunAttempt, j.ID).Return(j, nil).Maybe()
78 arts := mocks.NewCIArtifactStore(t)
79 arts.EXPECT().ArtifactUsage(mock.Anything, repo).Return(0, 1, nil).Maybe()
80 arts.EXPECT().GetArtifact(mock.Anything, a.ID).Return(a, nil).Maybe()
81 arts.EXPECT().ArtifactBlocks(mock.Anything, a.ID).Return(nil, nil).Maybe()
82 arts.EXPECT().ApplyArtifacts(mock.Anything, mock.Anything).Return(nil).Maybe()
83 clock := mocks.NewClock(t)
84 clock.EXPECT().Now().Return(time.Unix(1000, 0)).Maybe()
85 blobs, puts, _ := artifactBlobMock(t, false)
86 c := &CI{
87 cfg: CIConfig{
88 Store: store, Artifacts: arts, Blobs: blobs, Clock: clock,
89 ArtifactLimits: ports.ArtifactLimits{Quota: 1 << 30, MaxSize: 1 << 30},
90 },
91 uploadSlots: make(chan struct{}, slots),
92 }
93 return &slotFixture{c: c, secret: secret, job: j, puts: puts}
94}
95
96func (f *slotFixture) upload(ctx context.Context, r io.Reader) (int64, error) {
97 return f.c.UploadJobArtifactBlock(ctx, f.secret, f.job.ID, 7, []byte("blk"), r)
98}
99
100func TestArtifactUploadSlotsBound(t *testing.T) {
101 for _, mode := range []string{"success", "body error", "ctx cancel"} {
102 t.Run(mode, func(t *testing.T) {
103 f := newSlotFixture(t, 1)
104 ctx, cancel := context.WithCancel(t.Context())
105 defer cancel()
106 hold := &gateReader{ctx: ctx, started: make(chan struct{}), release: make(chan struct{}), data: []byte("held")}
107 if mode == "body error" {
108 hold.err = errors.New("connection reset")
109 }
110 type result struct {
111 n int64
112 err error
113 }
114 done := make(chan result, 1)
115 go func() { n, err := f.upload(ctx, hold); done <- result{n, err} }()
116 <-hold.started
117
118 // The slot is taken: the next upload is refused at once, before
119 // reading its body, and nothing is reserved or stored.
120 refused := &gateReader{ctx: t.Context(), data: []byte("x")}
121 _, err := f.upload(t.Context(), refused)
122 require.ErrorIs(t, err, domain.ErrResourceExhausted)
123 require.Zero(t, refused.reads)
124 require.Zero(t, *f.puts)
125 require.Zero(t, f.c.reservations.of(f.job.Repo).total)
126
127 switch mode {
128 case "ctx cancel":
129 cancel()
130 default:
131 close(hold.release)
132 }
133 res := <-done
134 if mode == "success" {
135 require.NoError(t, res.err)
136 require.EqualValues(t, 4, res.n)
137 } else {
138 require.Error(t, res.err)
139 }
140 require.Empty(t, f.c.uploadSlots, "the slot is released")
141
142 n, err := f.upload(t.Context(), &gateReader{ctx: t.Context(), data: []byte("again")})
143 require.NoError(t, err)
144 require.EqualValues(t, 5, n)
145 require.Empty(t, f.c.uploadSlots)
146 })
147 }
148}
149
150func TestArtifactAssemblyTakesASlot(t *testing.T) {
151 b, _, _ := artifactBlobMock(t, false)
152 c := &CI{cfg: CIConfig{Blobs: b}, uploadSlots: make(chan struct{}, 1)}
153 a := domain.CIArtifact{RunID: 1, ID: 2}
154 _, err := b.Put(t.Context(), a.BlockKey("a"), strings.NewReader("hi"), 2)
155 require.NoError(t, err)
156
157 // A full set of slots makes the assembly wait, not fail; it gives up
158 // only when its context ends.
159 c.uploadSlots <- struct{}{}
160 ctx, cancel := context.WithCancel(t.Context())
161 cancel()
162 _, _, err = c.assembleBlobs(ctx, a, []string{"a"}, 2)
163 require.ErrorIs(t, err, context.Canceled)
164
165 type result struct {
166 n int64
167 err error
168 }
169 done := make(chan result, 1)
170 go func() {
171 n, _, err := c.assembleBlobs(t.Context(), a, []string{"a"}, 2)
172 done <- result{n, err}
173 }()
174 <-c.uploadSlots // free the slot: the waiting assembly proceeds
175 res := <-done
176 require.NoError(t, res.err)
177 require.EqualValues(t, 2, res.n)
178 require.Empty(t, c.uploadSlots)
179}