fix(ci): bound concurrent artifact block uploads and reuse their buffers

commit: 8361f60249400dbacf799ec1a26de1ea55397d08
author: brice · 20h ago · 2026-10-21 16:28 UTC
committer: brice · 19h ago · 2026-10-21 17:53 UTC
parent: 6abe390
signature: GPG VerifiedCD88F9E564E0102646FB6077A1A07C7C7C3E349D · arch-gpg Verified for the committer email

3 files · +241 -6

M internal/usecase/ci.go +7 -4


@@ -249,9 +249,12 @@ func (c CIConfig) withDefaults() CIConfig {
249	249		// calls (a runner's update and an owner's cancel) never lose one another's
250	250		// write. Assignment takes the assignment lock before the lock of the run.
251	251		type CI struct {
252		-	cfg CIConfig
253		-	boot time.Time
254		-	queue chan ciTrigger
	252	+	cfg CIConfig
	253	+	// uploadSlots bounds the artifact uploads and assemblies buffering a
	254	+	// block in memory at once; nil (hand-built CI in tests) means unbounded.
	255	+	uploadSlots chan struct{}
	256	+	boot time.Time
	257	+	queue chan ciTrigger
255	258		// runEvents are the requested and completed events of runs, for the
256	259		// workflow_run workflows (ci_workflowrun.go).
257	260		runEvents chan runEvent

@@ -320,7 +323,7 @@ func NewCI(ctx context.Context, cfg CIConfig) *CI {
320	323		cfg = cfg.withDefaults()
321	324		run, stop := context.WithCancel(context.WithoutCancel(ctx))
322	325		return &CI{
323		-	pushes: map[domain.RepoName]int{}, pulls: map[pullKey]int{}, autoMerge: cfg.AutoMerge,
	326	+	pushes: map[domain.RepoName]int{}, uploadSlots: make(chan struct{}, ciArtifactUploadSlots), pulls: map[pullKey]int{}, autoMerge: cfg.AutoMerge,
324	327		cfg: cfg, boot: cfg.Clock.Now(), queue: make(chan ciTrigger, cfg.TriggerQueue), runEvents: make(chan runEvent, cfg.TriggerQueue), sched: newScheduler(cfg.Clock.Now()), pull: newPullQueue(cfg.TriggerQueue, cfg.PullRefreshes, cfg.PullAttempts),
325	328		ctx: run, stop: stop, flows: newWfCache(64), remoteSlots: make(chan struct{}, cfg.RemoteConcurrency), remoteCache: newRefsCache(),
326	329		heartbeats: map[string]time.Time{}, seen: map[string]time.Time{}, polls: map[string]pollRecord{},

 // calls (a runner's update and an owner's cancel) never lose one another's
 // write. Assignment takes the assignment lock before the lock of the run.
 type CI struct {
-	cfg   CIConfig
+	cfg CIConfig
-	boot  time.Time
+	// uploadSlots bounds the artifact uploads and assemblies buffering a
-	queue chan ciTrigger
+	// block in memory at once; nil (hand-built CI in tests) means unbounded.
+	uploadSlots chan struct{}
+	boot        time.Time
+	queue       chan ciTrigger
 	// runEvents are the requested and completed events of runs, for the
 	// workflow_run workflows (ci_workflowrun.go).
 	runEvents chan runEvent
 	cfg = cfg.withDefaults()
 	run, stop := context.WithCancel(context.WithoutCancel(ctx))
 	return &CI{
-		pushes: map[domain.RepoName]int{}, pulls: map[pullKey]int{}, autoMerge: cfg.AutoMerge,
+		pushes: map[domain.RepoName]int{}, uploadSlots: make(chan struct{}, ciArtifactUploadSlots), pulls: map[pullKey]int{}, autoMerge: cfg.AutoMerge,
 		cfg: cfg, boot: cfg.Clock.Now(), queue: make(chan ciTrigger, cfg.TriggerQueue), runEvents: make(chan runEvent, cfg.TriggerQueue), sched: newScheduler(cfg.Clock.Now()), pull: newPullQueue(cfg.TriggerQueue, cfg.PullRefreshes, cfg.PullAttempts),
 		ctx: run, stop: stop, flows: newWfCache(64), remoteSlots: make(chan struct{}, cfg.RemoteConcurrency), remoteCache: newRefsCache(),
 		heartbeats: map[string]time.Time{}, seen: map[string]time.Time{}, polls: map[string]pollRecord{},

M internal/usecase/ci_artifacts.go +55 -2


@@ -51,6 +51,44 @@ const (
51	51		commitTries = 5
52	52		)
53	53
	54	+	// ciArtifactUploadSlots bounds the artifact block uploads and assemblies that
	55	+	// hold an 8 MiB buffer at the same time, server-wide.
	56	+	const ciArtifactUploadSlots = 16
	57	+
	58	+	// blockBufs pools the 8 MiB buffers of uploads and assemblies. CIArtifactMaxBlockSize
	59	+	// and ports.BlobPartSize are both 8 MiB.
	60	+	var blockBufs = sync.Pool{New: func() any {
	61	+	b := make([]byte, max(domain.CIArtifactMaxBlockSize, ports.BlobPartSize))
	62	+	return &b
	63	+	}}
	64	+
	65	+	// acquireUploadSlot takes an upload slot without waiting; release frees it.
	66	+	func (c *CI) acquireUploadSlot() (release func(), err error) {
	67	+	if c.uploadSlots == nil {
	68	+	return func() {}, nil
	69	+	}
	70	+	select {
	71	+	case c.uploadSlots <- struct{}{}:
	72	+	return func() { <-c.uploadSlots }, nil
	73	+	default:
	74	+	return nil, fmt.Errorf("%w: too many artifact uploads in progress, retry", domain.ErrResourceExhausted)
	75	+	}
	76	+	}
	77	+
	78	+	// waitUploadSlot takes an upload slot, waiting for one until ctx ends: a
	79	+	// commit must not fail because block uploads are in flight.
	80	+	func (c *CI) waitUploadSlot(ctx context.Context) (release func(), err error) {
	81	+	if c.uploadSlots == nil {
	82	+	return func() {}, nil
	83	+	}
	84	+	select {
	85	+	case c.uploadSlots <- struct{}{}:
	86	+	return func() { <-c.uploadSlots }, nil
	87	+	case <-ctx.Done():
	88	+	return nil, ctx.Err()
	89	+	}
	90	+	}
	91	+
54	92		var (
55	93		// errNoArtifactWrite ends updateArtifacts without a write.
56	94		errNoArtifactWrite = errors.New("no write")

@@ -467,6 +505,12 @@ func (c *CI) UploadJobArtifactBlock(ctx context.Context, secret, jobID string, a
467	505		if a.State != domain.CIArtifactPending {
468	506		return 0, errCommitted(a.ID)
469	507		}
	508	+	// The slot is taken before anything is reserved, read or allocated.
	509	+	freeSlot, err := c.acquireUploadSlot()
	510	+	if err != nil {
	511	+	return 0, err
	512	+	}
	513	+	defer freeSlot()
470	514		// Every byte is reserved before it is written, against the charges and
471	515		// the other writes in flight (a block it replaces counts until the record
472	516		// drops it): nothing beyond the limits ever reaches the disk.

@@ -1070,7 +1114,9 @@ func (c *CI) sweepArtifactOrphans(ctx context.Context) {
1070	1114		func (c *CI) putUnsized(ctx context.Context, key string, r io.Reader) (int64, error) {
1071	1115		// A fixed backing array prevents bytes.Buffer's geometric growth from
1072	1116		// allocating beyond the block bound. Probe the next byte separately.
1073		-	buf := make([]byte, domain.CIArtifactMaxBlockSize)
	1117	+	bp := blockBufs.Get().(*[]byte)
	1118	+	defer blockBufs.Put(bp)
	1119	+	buf := (*bp)[:domain.CIArtifactMaxBlockSize]
1074	1120		n, err := io.ReadFull(r, buf)
1075	1121		if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
1076	1122		return 0, err

@@ -1101,6 +1147,11 @@ func (c *CI) removeBlobBlocks(ctx context.Context, a domain.CIArtifact) error {
1101	1147		// assembleBlobs keeps one 8 MiB part in memory, even when a block crosses
1102	1148		// part boundaries or is repeated in the commit list.
1103	1149		func (c *CI) assembleBlobs(ctx context.Context, a domain.CIArtifact, files []string, total int64) (int64, string, error) {
	1150	+	freeSlot, err := c.waitUploadSlot(ctx)
	1151	+	if err != nil {
	1152	+	return 0, "", err
	1153	+	}
	1154	+	defer freeSlot()
1104	1155		var upload ports.MultipartUpload
1105	1156		if total > ports.BlobPartSize {
1106	1157		var err error

@@ -1117,7 +1168,9 @@ func (c *CI) assembleBlobs(ctx context.Context, a domain.CIArtifact, files []str
1117	1168		}
1118	1169		}
1119	1170		}()
1120		-	buf := make([]byte, ports.BlobPartSize)
	1171	+	bp := blockBufs.Get().(*[]byte)
	1172	+	defer blockBufs.Put(bp)
	1173	+	buf := (*bp)[:ports.BlobPartSize]
1121	1174		used, part := 0, 1
1122	1175		var size int64
1123	1176		h := sha256.New()


@@ -51,6 +51,44 @@ const (
51		commitTries = 5	51		commitTries = 5
52		)	52		)
53			53
			54	+	// ciArtifactUploadSlots bounds the artifact block uploads and assemblies that
			55	+	// hold an 8 MiB buffer at the same time, server-wide.
			56	+	const ciArtifactUploadSlots = 16
			57	+
			58	+	// blockBufs pools the 8 MiB buffers of uploads and assemblies. CIArtifactMaxBlockSize
			59	+	// and ports.BlobPartSize are both 8 MiB.
			60	+	var blockBufs = sync.Pool{New: func() any {
			61	+	b := make([]byte, max(domain.CIArtifactMaxBlockSize, ports.BlobPartSize))
			62	+	return &b
			63	+	}}
			64	+
			65	+	// acquireUploadSlot takes an upload slot without waiting; release frees it.
			66	+	func (c *CI) acquireUploadSlot() (release func(), err error) {
			67	+	if c.uploadSlots == nil {
			68	+	return func() {}, nil
			69	+	}
			70	+	select {
			71	+	case c.uploadSlots <- struct{}{}:
			72	+	return func() { <-c.uploadSlots }, nil
			73	+	default:
			74	+	return nil, fmt.Errorf("%w: too many artifact uploads in progress, retry", domain.ErrResourceExhausted)
			75	+	}
			76	+	}
			77	+
			78	+	// waitUploadSlot takes an upload slot, waiting for one until ctx ends: a
			79	+	// commit must not fail because block uploads are in flight.
			80	+	func (c *CI) waitUploadSlot(ctx context.Context) (release func(), err error) {
			81	+	if c.uploadSlots == nil {
			82	+	return func() {}, nil
			83	+	}
			84	+	select {
			85	+	case c.uploadSlots <- struct{}{}:
			86	+	return func() { <-c.uploadSlots }, nil
			87	+	case <-ctx.Done():
			88	+	return nil, ctx.Err()
			89	+	}
			90	+	}
			91	+
54		var (	92		var (
55		// errNoArtifactWrite ends updateArtifacts without a write.	93		// errNoArtifactWrite ends updateArtifacts without a write.
56		errNoArtifactWrite = errors.New("no write")	94		errNoArtifactWrite = errors.New("no write")

@@ -467,6 +505,12 @@ func (c *CI) UploadJobArtifactBlock(ctx context.Context, secret, jobID string, a
467		if a.State != domain.CIArtifactPending {	505		if a.State != domain.CIArtifactPending {
468		return 0, errCommitted(a.ID)	506		return 0, errCommitted(a.ID)
469		}	507		}
			508	+	// The slot is taken before anything is reserved, read or allocated.
			509	+	freeSlot, err := c.acquireUploadSlot()
			510	+	if err != nil {
			511	+	return 0, err
			512	+	}
			513	+	defer freeSlot()
470		// Every byte is reserved before it is written, against the charges and	514		// Every byte is reserved before it is written, against the charges and
471		// the other writes in flight (a block it replaces counts until the record	515		// the other writes in flight (a block it replaces counts until the record
472		// drops it): nothing beyond the limits ever reaches the disk.	516		// drops it): nothing beyond the limits ever reaches the disk.

@@ -1070,7 +1114,9 @@ func (c *CI) sweepArtifactOrphans(ctx context.Context) {
1070		func (c *CI) putUnsized(ctx context.Context, key string, r io.Reader) (int64, error) {	1114		func (c *CI) putUnsized(ctx context.Context, key string, r io.Reader) (int64, error) {
1071		// A fixed backing array prevents bytes.Buffer's geometric growth from	1115		// A fixed backing array prevents bytes.Buffer's geometric growth from
1072		// allocating beyond the block bound. Probe the next byte separately.	1116		// allocating beyond the block bound. Probe the next byte separately.
1073	-	buf := make([]byte, domain.CIArtifactMaxBlockSize)	1117	+	bp := blockBufs.Get().(*[]byte)
			1118	+	defer blockBufs.Put(bp)
			1119	+	buf := (*bp)[:domain.CIArtifactMaxBlockSize]
1074		n, err := io.ReadFull(r, buf)	1120		n, err := io.ReadFull(r, buf)
1075		if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {	1121		if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
1076		return 0, err	1122		return 0, err

@@ -1101,6 +1147,11 @@ func (c *CI) removeBlobBlocks(ctx context.Context, a domain.CIArtifact) error {
1101		// assembleBlobs keeps one 8 MiB part in memory, even when a block crosses	1147		// assembleBlobs keeps one 8 MiB part in memory, even when a block crosses
1102		// part boundaries or is repeated in the commit list.	1148		// part boundaries or is repeated in the commit list.
1103		func (c *CI) assembleBlobs(ctx context.Context, a domain.CIArtifact, files []string, total int64) (int64, string, error) {	1149		func (c *CI) assembleBlobs(ctx context.Context, a domain.CIArtifact, files []string, total int64) (int64, string, error) {
			1150	+	freeSlot, err := c.waitUploadSlot(ctx)
			1151	+	if err != nil {
			1152	+	return 0, "", err
			1153	+	}
			1154	+	defer freeSlot()
1104		var upload ports.MultipartUpload	1155		var upload ports.MultipartUpload
1105		if total > ports.BlobPartSize {	1156		if total > ports.BlobPartSize {
1106		var err error	1157		var err error

@@ -1117,7 +1168,9 @@ func (c *CI) assembleBlobs(ctx context.Context, a domain.CIArtifact, files []str
1117		}	1168		}
1118		}	1169		}
1119		}()	1170		}()
1120	-	buf := make([]byte, ports.BlobPartSize)	1171	+	bp := blockBufs.Get().(*[]byte)
			1172	+	defer blockBufs.Put(bp)
			1173	+	buf := (*bp)[:ports.BlobPartSize]
1121		used, part := 0, 1	1174		used, part := 0, 1
1122		var size int64	1175		var size int64
1123		h := sha256.New()	1176		h := sha256.New()

A internal/usecase/ci_artifacts_slots_test.go +179 -0

@@ -0,0 +1,179 @@
	1	+	// SPDX-License-Identifier: AGPL-3.0-or-later
	2	+
	3	+	package usecase
	4	+
	5	+	import (
	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).
	25	+	type 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	+
	36	+	func (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	+
	59	+	type slotFixture struct {
	60	+	c *CI
	61	+	secret string
	62	+	job domain.CIJob
	63	+	puts *int
	64	+	}
	65	+
	66	+	func 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	+
	96	+	func (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	+
	100	+	func 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	+
	150	+	func 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	+	}

@@ -0,0 +1,179 @@
	1	+	// SPDX-License-Identifier: AGPL-3.0-or-later
	2	+
	3	+	package usecase
	4	+
	5	+	import (
	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).
	25	+	type 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	+
	36	+	func (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	+
	59	+	type slotFixture struct {
	60	+	c *CI
	61	+	secret string
	62	+	job domain.CIJob
	63	+	puts *int
	64	+	}
	65	+
	66	+	func 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	+
	96	+	func (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	+
	100	+	func 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	+
	150	+	func 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	+	}