1// SPDX-License-Identifier: AGPL-3.0-or-later
2
3package usecase
4
5import (
6 "bytes"
7 "io"
8 "testing"
9 "testing/iotest"
10
11 "github.com/stretchr/testify/assert"
12 "github.com/stretchr/testify/mock"
13 "github.com/stretchr/testify/require"
14
15 "bnema.dev/ermite/internal/domain"
16 "bnema.dev/ermite/internal/ports"
17 "bnema.dev/ermite/internal/ports/mocks"
18)
19
20// smallReads hands out at most 32 KiB per read, like io.Copy's buffer.
21type smallReads struct{ r io.Reader }
22
23func (s smallReads) Read(p []byte) (int, error) { return s.r.Read(p[:min(len(p), 32<<10)]) }
24
25func reserveFixture(t *testing.T, limits ports.ArtifactLimits) (*CI, *mocks.CIArtifactStore, domain.CIArtifact) {
26 t.Helper()
27 repo, err := domain.ParseRepoName("proj")
28 require.NoError(t, err)
29 a := domain.CIArtifact{ID: 7, Repo: repo, RunID: 1, State: domain.CIArtifactPending}
30 store := mocks.NewCIArtifactStore(t)
31 store.EXPECT().ArtifactUsage(mock.Anything, repo).Return(0, 1, nil).Maybe()
32 store.EXPECT().GetArtifact(mock.Anything, a.ID).Return(a, nil).Maybe()
33 return &CI{cfg: CIConfig{Artifacts: store, ArtifactLimits: limits}}, store, a
34}
35
36// A large block takes one reservation per MiB, not one per read.
37func TestReservationsPerMiB(t *testing.T) {
38 c, store, a := reserveFixture(t, ports.ArtifactLimits{Quota: 1 << 30, MaxSize: 1 << 30})
39 rr := &reservingReader{ctx: t.Context(), c: c, a: a, r: smallReads{bytes.NewReader(make([]byte, 8<<20))}}
40 n, err := io.Copy(io.Discard, rr)
41 require.NoError(t, err)
42 assert.EqualValues(t, 8<<20, n)
43 calls := 0
44 for _, call := range store.Calls {
45 if call.Method == "ArtifactUsage" {
46 calls++
47 }
48 }
49 assert.LessOrEqual(t, calls, 8, "one reservation per MiB for 256 reads")
50 assert.EqualValues(t, 8<<20, rr.reserved, "never more reserved than needed at the end of a whole MiB")
51
52 c.release(a, rr.reserved)
53 assert.Zero(t, c.reservations.of(a.Repo).total, "all given back")
54}
55
56// Near the limit a read reserves just what it needs: an upload that fits is
57// never refused because of the reservation ahead.
58func TestReservationsNearTheLimit(t *testing.T) {
59 const limit = 1<<20 + 100
60 c, _, a := reserveFixture(t, ports.ArtifactLimits{Quota: 1 << 30, MaxSize: limit})
61 rr := &reservingReader{ctx: t.Context(), c: c, a: a, r: iotest.HalfReader(bytes.NewReader(make([]byte, limit)))}
62 _, err := io.Copy(io.Discard, rr)
63 require.NoError(t, err)
64 assert.LessOrEqual(t, rr.reserved, int64(limit))
65
66 // One byte more than the limit: the read that crosses it fails and is not
67 // passed on.
68 c.release(a, rr.reserved)
69 rr = &reservingReader{ctx: t.Context(), c: c, a: a, r: bytes.NewReader(make([]byte, limit+1))}
70 n, err := io.Copy(io.Discard, rr)
71 require.ErrorIs(t, err, domain.ErrResourceExhausted)
72 assert.LessOrEqual(t, n, int64(limit))
73}
74
75// Repositories do not share a lock or a quota.
76func TestReservationsArePerRepository(t *testing.T) {
77 c, _, a := reserveFixture(t, ports.ArtifactLimits{Quota: 1 << 20, MaxSize: 1 << 20})
78 require.NoError(t, c.reserve(t.Context(), a, 1<<20))
79 assert.ErrorIs(t, c.reserve(t.Context(), a, 1), domain.ErrResourceExhausted)
80
81 other, err := domain.ParseRepoName("other")
82 require.NoError(t, err)
83 b := domain.CIArtifact{ID: 8, Repo: other, RunID: 2, State: domain.CIArtifactPending}
84 store := c.cfg.Artifacts.(*mocks.CIArtifactStore)
85 store.EXPECT().ArtifactUsage(mock.Anything, other).Return(0, 1, nil)
86 store.EXPECT().GetArtifact(mock.Anything, b.ID).Return(b, nil)
87 require.NoError(t, c.reserve(t.Context(), b, 1<<20), "the quota of another repository is its own")
88 assert.NotSame(t, c.reservations.of(a.Repo), c.reservations.of(b.Repo))
89}