Cancel remaining index fetches on first error and test parallel_fetch wiring

internal/api/v2/virtuals_test.go

@@ -0,0 +1,56 @@
+package v2
+
+import (
+	"encoding/json"
+	"net/http/httptest"
+	"testing"
+
+	"git.unkin.net/unkin/artifactapi/internal/database"
+)
+
+func TestVirtualParallelFetchRoundTrip(t *testing.T) {
+	if testDSN == "" {
+		t.Skip("Docker unavailable")
+	}
+	db, err := database.New(testDSN)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+	h := NewVirtualsHandler(db).Routes()
+	parallel := func() bool {
+		t.Helper()
+		w := httptest.NewRecorder()
+		h.ServeHTTP(w, httptest.NewRequest("GET", "/vpf", nil))
+		var got map[string]any
+		if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil || w.Code != 200 {
+			t.Fatalf("get = %d %s", w.Code, w.Body)
+		}
+		v, ok := got["parallel_fetch"].(bool)
+		if !ok {
+			t.Fatalf("parallel_fetch missing from %s", w.Body)
+		}
+		return v
+	}
+	const body = `"package_type":"rpm","members":["a"]`
+
+	steps := []struct {
+		method, path, json string
+		want               bool
+	}{
+		{"POST", "/", `{"name":"vpf",` + body + `}`, false},
+		{"PUT", "/vpf", `{` + body + `,"parallel_fetch":true}`, true},
+		{"PUT", "/vpf", `{` + body + `,"parallel_fetch":false}`, false},
+	}
+	for i, s := range steps {
+		if c := do(t, h, s.method, s.path, s.json); c/100 != 2 {
+			t.Fatalf("step %d %s %s = %d", i, s.method, s.path, c)
+		}
+		if got := parallel(); got != s.want {
+			t.Errorf("step %d: parallel_fetch = %v, want %v", i, got, s.want)
+		}
+	}
+	if c := do(t, h, "DELETE", "/vpf", ""); c != 204 {
+		t.Fatalf("delete = %d", c)
+	}
+}

internal/database/database_test.go

@@ -335,6 +335,40 @@ func TestDeleteRemoteDropsLocalFiles(t *testing.T) {
 	}
 }
 
+func TestVirtualParallelFetchRoundTrip(t *testing.T) {
+	requireDB(t)
+	if err := testDB.CreateVirtual(ctx(), &models.Virtual{
+		Name: "v-parallel", PackageType: models.PackageRPM, Members: []string{"a"}, ParallelFetch: true,
+	}); err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = testDB.DeleteVirtual(ctx(), "v-parallel") }()
+	v, err := testDB.GetVirtual(ctx(), "v-parallel")
+	if err != nil || !v.ParallelFetch {
+		t.Fatalf("created parallel_fetch = %v, %v; want true", v, err)
+	}
+	v.ParallelFetch = false
+	if err := testDB.UpdateVirtual(ctx(), v); err != nil {
+		t.Fatal(err)
+	}
+	vs, err := testDB.ListVirtuals(ctx())
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, l := range vs {
+		if l.Name == "v-parallel" && l.ParallelFetch {
+			t.Fatal("updated parallel_fetch not persisted")
+		}
+	}
+	if err := testDB.CreateVirtual(ctx(), &models.Virtual{Name: "v-serial", PackageType: models.PackageRPM, Members: []string{"a"}}); err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = testDB.DeleteVirtual(ctx(), "v-serial") }()
+	if v, err := testDB.GetVirtual(ctx(), "v-serial"); err != nil || v.ParallelFetch {
+		t.Fatalf("default parallel_fetch = %v, %v; want false", v, err)
+	}
+}
+
 func TestVirtualsCRUD(t *testing.T) {
 	requireDB(t)
 	if err := testDB.CreateVirtual(ctx(), &models.Virtual{

internal/virtual/engine.go

@@ -31,6 +31,7 @@ type Engine struct {
 	debMember   func(ctx context.Context, name, base string, parallel bool) (*DebMember, error)
 	apkMember   func(ctx context.Context, name, dir string) (*AlpineMember, error)
 	npmMember   func(ctx context.Context, name, pkg string) (*NPMMember, error)
+	memberPath  func(ctx context.Context, remote models.Remote, path string) ([]byte, error)
 
 	// NPMLocal returns a local npm repo's packument with tarball paths relative
 	// to the repo, or ErrNotFound.
@@ -92,6 +93,7 @@ func NewEngine(db *database.DB, proxyEngine *proxy.Engine) *Engine {
 	e.ansibleDoc = e.ansibleMemberDoc
 	e.goMember = e.fetchGoMember
 	e.puppetDoc = e.fetchMemberPath
+	e.memberPath = e.fetchMemberPath
 	return e
 }
 
@@ -436,11 +438,11 @@ func (e *Engine) fetchDebMember(ctx context.Context, name, base string, parallel
 	if err != nil {
 		return nil, fmt.Errorf("remote %q: %w", name, err)
 	}
-	rel, err := e.fetchMemberPath(ctx, *remote, base+"Release")
+	rel, err := e.memberPath(ctx, *remote, base+"Release")
 	if err != nil {
 		return nil, err
 	}
-	pkgs, err := fetchFiles(parallel, debIndexFiles(parseDebRelease(rel)), func(file string) ([]byte, error) {
+	pkgs, err := fetchFiles(ctx, parallel, debIndexFiles(parseDebRelease(rel)), func(ctx context.Context, file string) ([]byte, error) {
 		return e.fetchDecompressed(ctx, *remote, base, file)
 	})
 	if err != nil {
@@ -561,7 +563,7 @@ func (e *Engine) fetchRPMMember(ctx context.Context, name string, parallel bool)
 	if err != nil {
 		return nil, fmt.Errorf("remote %q: %w", name, err)
 	}
-	repomd, err := e.fetchMemberPath(ctx, *remote, "repodata/repomd.xml")
+	repomd, err := e.memberPath(ctx, *remote, "repodata/repomd.xml")
 	if err != nil {
 		return nil, err
 	}
@@ -569,7 +571,7 @@ func (e *Engine) fetchRPMMember(ctx context.Context, name string, parallel bool)
 	if err != nil {
 		return nil, err
 	}
-	data, err := fetchFiles(parallel, locs, func(href string) ([]byte, error) {
+	data, err := fetchFiles(ctx, parallel, locs, func(ctx context.Context, href string) ([]byte, error) {
 		return e.fetchDecompressed(ctx, *remote, "", href)
 	})
 	if err != nil {
@@ -579,7 +581,7 @@ func (e *Engine) fetchRPMMember(ctx context.Context, name string, parallel bool)
 }
 
 func (e *Engine) fetchDecompressed(ctx context.Context, remote models.Remote, base, file string) ([]byte, error) {
-	raw, err := e.fetchMemberPath(ctx, remote, base+file)
+	raw, err := e.memberPath(ctx, remote, base+file)
 	if err != nil {
 		return nil, err
 	}
@@ -594,12 +596,12 @@ func (e *Engine) fetchDecompressed(ctx context.Context, remote models.Remote, ba
 const maxFileFetches = 4
 
 // fetchFiles fetches each key's file, one at a time or up to maxFileFetches at
-// once, failing on the first error.
-func fetchFiles(parallel bool, files map[string]string, fetch func(file string) ([]byte, error)) (map[string][]byte, error) {
+// once, failing on the first error and cancelling the rest.
+func fetchFiles(ctx context.Context, parallel bool, files map[string]string, fetch func(ctx context.Context, file string) ([]byte, error)) (map[string][]byte, error) {
 	out := make(map[string][]byte, len(files))
 	if !parallel {
 		for k, f := range files {
-			body, err := fetch(f)
+			body, err := fetch(ctx, f)
 			if err != nil {
 				return nil, err
 			}
@@ -608,11 +610,14 @@ func fetchFiles(parallel bool, files map[string]string, fetch func(file string)
 		return out, nil
 	}
 	var mu sync.Mutex
-	var g errgroup.Group
+	g, gctx := errgroup.WithContext(ctx)
 	g.SetLimit(maxFileFetches)
 	for k, f := range files {
 		g.Go(func() error {
-			body, err := fetch(f)
+			if err := gctx.Err(); err != nil {
+				return err
+			}
+			body, err := fetch(gctx, f)
 			if err != nil {
 				return err
 			}

internal/virtual/engine_test.go

@@ -389,7 +389,7 @@ func TestFetchFilesParallelRunsConcurrently(t *testing.T) {
 		}
 		close(release)
 	}()
-	got, err := fetchFiles(true, fetchFilesInput(6), func(f string) ([]byte, error) {
+	got, err := fetchFiles(context.Background(), true, fetchFilesInput(6), func(_ context.Context, f string) ([]byte, error) {
 		n := inFlight.Add(1)
 		defer inFlight.Add(-1)
 		for p := peak.Load(); n > p && !peak.CompareAndSwap(p, n); p = peak.Load() {
@@ -416,7 +416,7 @@ func TestFetchFilesParallelRunsConcurrently(t *testing.T) {
 
 func TestFetchFilesSerialByDefault(t *testing.T) {
 	var inFlight atomic.Int32
-	_, err := fetchFiles(false, fetchFilesInput(4), func(f string) ([]byte, error) {
+	_, err := fetchFiles(context.Background(), false, fetchFilesInput(4), func(_ context.Context, f string) ([]byte, error) {
 		defer inFlight.Add(-1)
 		if inFlight.Add(1) > 1 {
 			return nil, errors.New("concurrent fetch")
@@ -432,7 +432,7 @@ func TestFetchFilesSerialByDefault(t *testing.T) {
 func TestFetchFilesErrorFailsMember(t *testing.T) {
 	boom := errors.New("boom")
 	for _, parallel := range []bool{false, true} {
-		got, err := fetchFiles(parallel, fetchFilesInput(5), func(f string) ([]byte, error) {
+		got, err := fetchFiles(context.Background(), parallel, fetchFilesInput(5), func(_ context.Context, f string) ([]byte, error) {
 			if f == "f2" {
 				return nil, boom
 			}
@@ -469,3 +469,98 @@ func TestMergeRPMPassesParallelFetchAndKeepsMemberOrder(t *testing.T) {
 		t.Fatalf("member priority order lost:\n%s", body)
 	}
 }
+
+func TestFetchFilesErrorCancelsRest(t *testing.T) {
+	boom := errors.New("boom")
+	var calls atomic.Int32
+	_, err := fetchFiles(context.Background(), true, fetchFilesInput(8), func(ctx context.Context, _ string) ([]byte, error) {
+		if calls.Add(1) == 1 {
+			return nil, boom
+		}
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-time.After(5 * time.Second):
+			return nil, errors.New("in-flight fetch not cancelled")
+		}
+	})
+	if !errors.Is(err, boom) {
+		t.Fatalf("got %v, want boom", err)
+	}
+	if n := calls.Load(); n > maxFileFetches {
+		t.Fatalf("%d fetches started after the first error, want at most %d", n, maxFileFetches)
+	}
+}
+
+// overlapEngine serves fixed member index paths, holding each index file fetch
+// until all are in flight or a short wait expires, and records the peak.
+func overlapEngine(index map[string][]byte, indexFiles int, peak *atomic.Int32) *Engine {
+	e := fakeEngine(nil)
+	var inFlight atomic.Int32
+	e.memberPath = func(_ context.Context, _ models.Remote, path string) ([]byte, error) {
+		body, ok := index[path]
+		if !ok {
+			return nil, fmt.Errorf("%s: %w", path, ErrNotFound)
+		}
+		if strings.HasSuffix(path, "Release") || strings.HasSuffix(path, "repomd.xml") {
+			return body, nil
+		}
+		n := inFlight.Add(1)
+		defer inFlight.Add(-1)
+		for p := peak.Load(); n > p && !peak.CompareAndSwap(p, n); p = peak.Load() {
+		}
+		for deadline := time.Now().Add(50 * time.Millisecond); peak.Load() < int32(indexFiles) && time.Now().Before(deadline); {
+			time.Sleep(time.Millisecond)
+		}
+		return body, nil
+	}
+	return e
+}
+
+func TestFetchRPMMemberHonoursParallelFetch(t *testing.T) {
+	index := map[string][]byte{
+		"repodata/repomd.xml": []byte(`<repomd>` +
+			`<data type="primary"><location href="repodata/p.xml"/></data>` +
+			`<data type="filelists"><location href="repodata/f.xml"/></data>` +
+			`<data type="other"><location href="repodata/o.xml"/></data></repomd>`),
+		"repodata/p.xml": []byte("P"), "repodata/f.xml": []byte("F"), "repodata/o.xml": []byte("O"),
+	}
+	for parallel, want := range map[bool]int32{false: 1, true: 3} {
+		var peak atomic.Int32
+		m, err := overlapEngine(index, 3, &peak).fetchRPMMember(context.Background(), "a", parallel)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if peak.Load() != want {
+			t.Errorf("parallel=%v: peak concurrency %d, want %d", parallel, peak.Load(), want)
+		}
+		if string(m.Data["primary"]) != "P" || string(m.Data["filelists"]) != "F" || string(m.Data["other"]) != "O" {
+			t.Errorf("parallel=%v: data %q", parallel, m.Data)
+		}
+	}
+}
+
+func TestFetchDebMemberHonoursParallelFetch(t *testing.T) {
+	index := map[string][]byte{
+		"dists/s/Release": []byte("Suite: s\nSHA256:\n" +
+			" aa 1 main/binary-amd64/Packages\n" +
+			" bb 1 contrib/binary-amd64/Packages\n" +
+			" cc 1 non-free/binary-amd64/Packages\n"),
+		"dists/s/main/binary-amd64/Packages":     []byte("M"),
+		"dists/s/contrib/binary-amd64/Packages":  []byte("C"),
+		"dists/s/non-free/binary-amd64/Packages": []byte("N"),
+	}
+	for parallel, want := range map[bool]int32{false: 1, true: 3} {
+		var peak atomic.Int32
+		m, err := overlapEngine(index, 3, &peak).fetchDebMember(context.Background(), "a", "dists/s/", parallel)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if peak.Load() != want {
+			t.Errorf("parallel=%v: peak concurrency %d, want %d", parallel, peak.Load(), want)
+		}
+		if string(m.Packages["main/binary-amd64/"]) != "M" || string(m.Packages["non-free/binary-amd64/"]) != "N" {
+			t.Errorf("parallel=%v: packages %q", parallel, m.Packages)
+		}
+	}
+}