Add parallel_fetch toggle for virtual member index files

2026-10-25 13:09:28 +11:00
parent 9294686ea6
commit 56be0ac266
5 changed files with 196 additions and 49 deletions
@@ -8,12 +8,12 @@ import (

 func (db *DB) GetVirtual(ctx context.Context, name string) (*models.Virtual, error) {
 	row := db.Pool.QueryRow(ctx, `
-		SELECT name, package_type, description, members, managed_by, created_at, updated_at
+		SELECT name, package_type, description, members, managed_by, parallel_fetch, created_at, updated_at
 		FROM virtuals WHERE name = $1
 	`, name)

 	var v models.Virtual
-	err := row.Scan(&v.Name, &v.PackageType, &v.Description, &v.Members, &v.ManagedBy, &v.CreatedAt, &v.UpdatedAt)
+	err := row.Scan(&v.Name, &v.PackageType, &v.Description, &v.Members, &v.ManagedBy, &v.ParallelFetch, &v.CreatedAt, &v.UpdatedAt)
 	if err != nil {
 		return nil, err
 	}
@@ -22,7 +22,7 @@ func (db *DB) GetVirtual(ctx context.Context, name string) (*models.Virtual, err

 func (db *DB) ListVirtuals(ctx context.Context) ([]models.Virtual, error) {
 	rows, err := db.Pool.Query(ctx, `
-		SELECT name, package_type, description, members, managed_by, created_at, updated_at
+		SELECT name, package_type, description, members, managed_by, parallel_fetch, created_at, updated_at
 		FROM virtuals ORDER BY name
 	`)
 	if err != nil {
@@ -33,7 +33,7 @@ func (db *DB) ListVirtuals(ctx context.Context) ([]models.Virtual, error) {
 	var virtuals []models.Virtual
 	for rows.Next() {
 		var v models.Virtual
-		if err := rows.Scan(&v.Name, &v.PackageType, &v.Description, &v.Members, &v.ManagedBy, &v.CreatedAt, &v.UpdatedAt); err != nil {
+		if err := rows.Scan(&v.Name, &v.PackageType, &v.Description, &v.Members, &v.ManagedBy, &v.ParallelFetch, &v.CreatedAt, &v.UpdatedAt); err != nil {
 			return nil, err
 		}
 		virtuals = append(virtuals, v)
@@ -43,18 +43,18 @@ func (db *DB) ListVirtuals(ctx context.Context) ([]models.Virtual, error) {

 func (db *DB) CreateVirtual(ctx context.Context, v *models.Virtual) error {
 	_, err := db.Pool.Exec(ctx, `
-		INSERT INTO virtuals (name, package_type, description, members, managed_by)
-		VALUES ($1, $2, $3, $4, $5)
-	`, v.Name, v.PackageType, v.Description, v.Members, v.ManagedBy)
+		INSERT INTO virtuals (name, package_type, description, members, managed_by, parallel_fetch)
+		VALUES ($1, $2, $3, $4, $5, $6)
+	`, v.Name, v.PackageType, v.Description, v.Members, v.ManagedBy, v.ParallelFetch)
 	return err
 }

 func (db *DB) UpdateVirtual(ctx context.Context, v *models.Virtual) error {
 	_, err := db.Pool.Exec(ctx, `
 		UPDATE virtuals SET
-			package_type=$2, description=$3, members=$4, managed_by=$5, updated_at=NOW()
+			package_type=$2, description=$3, members=$4, managed_by=$5, parallel_fetch=$6, updated_at=NOW()
 		WHERE name=$1
-	`, v.Name, v.PackageType, v.Description, v.Members, v.ManagedBy)
+	`, v.Name, v.PackageType, v.Description, v.Members, v.ManagedBy, v.ParallelFetch)
 	return err
 }

@@ -19,6 +19,7 @@ import (
 	"git.unkin.net/unkin/artifactapi/internal/provider"
 	"git.unkin.net/unkin/artifactapi/internal/proxy"
 	"git.unkin.net/unkin/artifactapi/pkg/models"
+	"golang.org/x/sync/errgroup"
 	"golang.org/x/sync/singleflight"
 )

@@ -26,8 +27,8 @@ type Engine struct {
 	db          *database.DB
 	proxyEngine *proxy.Engine
 	getRemote   func(context.Context, string) (*models.Remote, error)
-	rpmMember   func(context.Context, string) (*RPMMember, error)
-	debMember   func(ctx context.Context, name, base string) (*DebMember, error)
+	rpmMember   func(ctx context.Context, name string, parallel bool) (*RPMMember, error)
+	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)

@@ -353,7 +354,7 @@ func (e *Engine) merged(ctx context.Context, key string, merge func(context.Cont
 }

 func (e *Engine) mergeRPM(ctx context.Context, virt models.Virtual) (*RPMRepo, error) {
-	members, err := fetchMembers(virt, func(name string) (*RPMMember, error) { return e.rpmMember(ctx, name) })
+	members, err := fetchMembers(virt, func(name string) (*RPMMember, error) { return e.rpmMember(ctx, name, virt.ParallelFetch) })
 	if err != nil {
 		return nil, err
 	}
@@ -402,7 +403,7 @@ func (e *Engine) fetchDeb(ctx context.Context, virt models.Virtual, reqPath stri
 	}

 	repo, err := e.merged(ctx, key, func(ctx context.Context) (mergedRepo, error) {
-		members, err := fetchMembers(virt, func(name string) (*DebMember, error) { return e.debMember(ctx, name, base) })
+		members, err := fetchMembers(virt, func(name string) (*DebMember, error) { return e.debMember(ctx, name, base, virt.ParallelFetch) })
 		if err != nil {
 			return nil, err
 		}
@@ -430,7 +431,7 @@ func debContentType(name string) string {

 // fetchDebMember reads a member's Release under base and each binary Packages
 // index it lists.
-func (e *Engine) fetchDebMember(ctx context.Context, name, base string) (*DebMember, error) {
+func (e *Engine) fetchDebMember(ctx context.Context, name, base string, parallel bool) (*DebMember, error) {
 	remote, err := e.getRemote(ctx, name)
 	if err != nil {
 		return nil, fmt.Errorf("remote %q: %w", name, err)
@@ -439,17 +440,13 @@ func (e *Engine) fetchDebMember(ctx context.Context, name, base string) (*DebMem
 	if err != nil {
 		return nil, err
 	}
-	m := &DebMember{RemoteName: name, Release: rel, Packages: map[string][]byte{}}
-	for dir, file := range debIndexFiles(parseDebRelease(rel)) {
-		raw, err := e.fetchMemberPath(ctx, *remote, base+file)
-		if err != nil {
-			return nil, err
-		}
-		if m.Packages[dir], err = decompress(file, raw); err != nil {
-			return nil, fmt.Errorf("decompress %s: %w", file, err)
-		}
+	pkgs, err := fetchFiles(parallel, debIndexFiles(parseDebRelease(rel)), func(file string) ([]byte, error) {
+		return e.fetchDecompressed(ctx, *remote, base, file)
+	})
+	if err != nil {
+		return nil, err
 	}
-	return m, nil
+	return &DebMember{RemoteName: name, Release: rel, Packages: pkgs}, nil
 }

 func (e *Engine) mergedAlpine(ctx context.Context, virt models.Virtual, dir string) (mergedRepo, error) {
@@ -559,7 +556,7 @@ func (e *Engine) cachedFile(key, path string) ([]byte, bool) {
 	return nil, false
 }

-func (e *Engine) fetchRPMMember(ctx context.Context, name string) (*RPMMember, error) {
+func (e *Engine) fetchRPMMember(ctx context.Context, name string, parallel bool) (*RPMMember, error) {
 	remote, err := e.getRemote(ctx, name)
 	if err != nil {
 		return nil, fmt.Errorf("remote %q: %w", name, err)
@@ -572,17 +569,63 @@ func (e *Engine) fetchRPMMember(ctx context.Context, name string) (*RPMMember, e
 	if err != nil {
 		return nil, err
 	}
-	m := &RPMMember{RemoteName: name, Bases: remote.UpstreamPool(), Timestamp: ts, Data: map[string][]byte{}}
-	for t, href := range locs {
-		raw, err := e.fetchMemberPath(ctx, *remote, href)
-		if err != nil {
-			return nil, err
-		}
-		if m.Data[t], err = decompress(href, raw); err != nil {
-			return nil, fmt.Errorf("decompress %s: %w", href, err)
-		}
+	data, err := fetchFiles(parallel, locs, func(href string) ([]byte, error) {
+		return e.fetchDecompressed(ctx, *remote, "", href)
+	})
+	if err != nil {
+		return nil, err
 	}
-	return m, nil
+	return &RPMMember{RemoteName: name, Bases: remote.UpstreamPool(), Timestamp: ts, Data: data}, nil
+}
+
+func (e *Engine) fetchDecompressed(ctx context.Context, remote models.Remote, base, file string) ([]byte, error) {
+	raw, err := e.fetchMemberPath(ctx, remote, base+file)
+	if err != nil {
+		return nil, err
+	}
+	body, err := decompress(file, raw)
+	if err != nil {
+		return nil, fmt.Errorf("decompress %s: %w", file, err)
+	}
+	return body, nil
+}
+
+// maxFileFetches bounds the concurrent index file fetches per member.
+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) {
+	out := make(map[string][]byte, len(files))
+	if !parallel {
+		for k, f := range files {
+			body, err := fetch(f)
+			if err != nil {
+				return nil, err
+			}
+			out[k] = body
+		}
+		return out, nil
+	}
+	var mu sync.Mutex
+	var g errgroup.Group
+	g.SetLimit(maxFileFetches)
+	for k, f := range files {
+		g.Go(func() error {
+			body, err := fetch(f)
+			if err != nil {
+				return err
+			}
+			mu.Lock()
+			out[k] = body
+			mu.Unlock()
+			return nil
+		})
+	}
+	if err := g.Wait(); err != nil {
+		return nil, err
+	}
+	return out, nil
 }

 // fetchMemberPath reads one path from a member the same way its own route
@@ -19,7 +19,7 @@ func fakeEngine(members map[string]*RPMMember) *Engine {
 		getRemote: func(_ context.Context, name string) (*models.Remote, error) {
 			return &models.Remote{Name: name, RepoType: models.RepoTypeRemote}, nil
 		},
-		rpmMember: func(_ context.Context, name string) (*RPMMember, error) {
+		rpmMember: func(_ context.Context, name string, _ bool) (*RPMMember, error) {
 			if m := members[name]; m != nil {
 				return m, nil
 			}
@@ -75,7 +75,7 @@ func TestFetchRPMLocalMemberWithoutRepodataIsUpstreamError(t *testing.T) {
 	e := fakeEngine(map[string]*RPMMember{
 		"a": {RemoteName: "a", Data: map[string][]byte{"primary": primaryXML(primaryPkgXML("foo", "1", "aaa", "foo.rpm"))}},
 	})
-	e.rpmMember = func(_ context.Context, name string) (*RPMMember, error) {
+	e.rpmMember = func(_ context.Context, name string, _ bool) (*RPMMember, error) {
 		if name == "local" {
 			return nil, fmt.Errorf("local/repodata/repomd.xml: %w", ErrNotFound)
 		}
@@ -91,7 +91,7 @@ func TestFetchRPMMergeIsShared(t *testing.T) {
 	var calls atomic.Int32
 	e := fakeEngine(nil)
 	e.mergeTTL = time.Minute
-	e.rpmMember = func(_ context.Context, name string) (*RPMMember, error) {
+	e.rpmMember = func(_ context.Context, name string, _ bool) (*RPMMember, error) {
 		calls.Add(1)
 		time.Sleep(20 * time.Millisecond)
 		return &RPMMember{RemoteName: name, Data: map[string][]byte{"primary": primaryXML(primaryPkgXML("foo", "1", "aaa", "foo.rpm"))}}, nil
@@ -160,7 +160,7 @@ func TestMemberRedirect(t *testing.T) {

 func debEngine(members map[string]*DebMember) *Engine {
 	e := fakeEngine(nil)
-	e.debMember = func(_ context.Context, name, base string) (*DebMember, error) {
+	e.debMember = func(_ context.Context, name, base string, _ bool) (*DebMember, error) {
 		if m := members[name]; m != nil && base == "" {
 			return m, nil
 		}
@@ -213,7 +213,7 @@ func TestFetchDebMergeIsSharedPerSuite(t *testing.T) {
 	var calls atomic.Int32
 	e := fakeEngine(nil)
 	e.mergeTTL = time.Minute
-	e.debMember = func(_ context.Context, name, base string) (*DebMember, error) {
+	e.debMember = func(_ context.Context, name, base string, _ bool) (*DebMember, error) {
 		calls.Add(1)
 		m := flatDebMember(name, "", debStanza("foo", "1", "amd64", "pool/foo.deb"))
 		return &m, nil
@@ -307,7 +307,9 @@ func suiteDebMember(name, version string) *DebMember {
 func TestFetchDebByHashSurvivesRemerge(t *testing.T) {
 	version := "1"
 	e := fakeEngine(nil)
-	e.debMember = func(_ context.Context, name, _ string) (*DebMember, error) { return suiteDebMember(name, version), nil }
+	e.debMember = func(_ context.Context, name, _ string, _ bool) (*DebMember, error) {
+		return suiteDebMember(name, version), nil
+	}
 	virt := debVirt("a")
 	fetch := func(p string) []byte {
 		body, _, err := e.Fetch(context.Background(), virt, p, "")
@@ -340,7 +342,9 @@ func TestFetchDebByHashSurvivesRemerge(t *testing.T) {
 func TestFetchDebBoundsSuitesPerVirtual(t *testing.T) {
 	e := fakeEngine(nil)
 	e.mergeTTL = time.Minute
-	e.debMember = func(_ context.Context, name, _ string) (*DebMember, error) { return suiteDebMember(name, "1"), nil }
+	e.debMember = func(_ context.Context, name, _ string, _ bool) (*DebMember, error) {
+		return suiteDebMember(name, "1"), nil
+	}
 	other := models.Virtual{Name: "other", PackageType: models.PackageDeb, Members: []string{"a"}}
 	if _, _, err := e.Fetch(context.Background(), other, "dists/keep/Release", ""); err != nil {
 		t.Fatal(err)
@@ -367,3 +371,101 @@ func TestFetchDebBoundsSuitesPerVirtual(t *testing.T) {
 		t.Fatal("another virtual's suites must not be evicted")
 	}
 }
+
+func fetchFilesInput(n int) map[string]string {
+	files := map[string]string{}
+	for i := range n {
+		files[fmt.Sprintf("k%d", i)] = fmt.Sprintf("f%d", i)
+	}
+	return files
+}
+
+func TestFetchFilesParallelRunsConcurrently(t *testing.T) {
+	var inFlight, peak atomic.Int32
+	release := make(chan struct{})
+	go func() {
+		for peak.Load() < maxFileFetches {
+			time.Sleep(time.Millisecond)
+		}
+		close(release)
+	}()
+	got, err := fetchFiles(true, fetchFilesInput(6), func(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() {
+		}
+		select {
+		case <-release:
+		case <-time.After(5 * time.Second):
+			return nil, errors.New("fetches never overlapped")
+		}
+		return []byte(f), nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if peak.Load() != maxFileFetches {
+		t.Fatalf("peak concurrency %d, want %d", peak.Load(), maxFileFetches)
+	}
+	for k, f := range fetchFilesInput(6) {
+		if string(got[k]) != f {
+			t.Fatalf("%s = %q, want %q", k, got[k], f)
+		}
+	}
+}
+
+func TestFetchFilesSerialByDefault(t *testing.T) {
+	var inFlight atomic.Int32
+	_, err := fetchFiles(false, fetchFilesInput(4), func(f string) ([]byte, error) {
+		defer inFlight.Add(-1)
+		if inFlight.Add(1) > 1 {
+			return nil, errors.New("concurrent fetch")
+		}
+		time.Sleep(time.Millisecond)
+		return []byte(f), nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+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) {
+			if f == "f2" {
+				return nil, boom
+			}
+			return []byte(f), nil
+		})
+		if !errors.Is(err, boom) || got != nil {
+			t.Fatalf("parallel=%v: got %v, %v; want boom and no files", parallel, got, err)
+		}
+	}
+}
+
+func TestMergeRPMPassesParallelFetchAndKeepsMemberOrder(t *testing.T) {
+	e := fakeEngine(nil)
+	var sawParallel atomic.Bool
+	e.rpmMember = func(_ context.Context, name string, parallel bool) (*RPMMember, error) {
+		sawParallel.Store(parallel)
+		if name == "a" {
+			time.Sleep(20 * time.Millisecond)
+		}
+		return &RPMMember{RemoteName: name, Data: map[string][]byte{"primary": primaryXML(primaryPkgXML("foo", "1", name+"sum", name+".rpm"))}}, nil
+	}
+	virt := rpmVirt("a", "b")
+	virt.ParallelFetch = true
+	repo, err := e.mergeRPM(context.Background(), virt)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !sawParallel.Load() {
+		t.Fatal("parallel_fetch not passed to member fetch")
+	}
+	primary := regexp.MustCompile(`repodata/[0-9a-f]+-primary\.xml\.gz`).FindString(string(repo.Repomd))
+	body := string(gunzip(t, repo.Files[primary]))
+	if a, b := strings.Index(body, "a.rpm"), strings.Index(body, "b.rpm"); a < 0 || b >= 0 && b < a {
+		t.Fatalf("member priority order lost:\n%s", body)
+	}
+}
@@ -0,0 +1 @@
+ALTER TABLE virtuals ADD COLUMN IF NOT EXISTS parallel_fetch BOOLEAN NOT NULL DEFAULT FALSE;
@@ -3,11 +3,12 @@ package models
 import "time"

 type Virtual struct {
-	Name        string      `json:"name"`
-	PackageType PackageType `json:"package_type"`
-	Description string      `json:"description,omitempty"`
-	Members     []string    `json:"members"`
-	ManagedBy   string      `json:"managed_by,omitempty"`
-	CreatedAt   time.Time   `json:"created_at"`
-	UpdatedAt   time.Time   `json:"updated_at"`
+	Name          string      `json:"name"`
+	PackageType   PackageType `json:"package_type"`
+	Description   string      `json:"description,omitempty"`
+	Members       []string    `json:"members"`
+	ManagedBy     string      `json:"managed_by,omitempty"`
+	ParallelFetch bool        `json:"parallel_fetch"`
+	CreatedAt     time.Time   `json:"created_at"`
+	UpdatedAt     time.Time   `json:"updated_at"`
 }
				`@@ -0,0 +1 @@`
				`ALTER TABLE virtuals ADD COLUMN IF NOT EXISTS parallel_fetch BOOLEAN NOT NULL DEFAULT FALSE;`