Refresh the Julia name map only from the advertised registry

internal/handler/julia.go

@@ -12,6 +12,7 @@ import (
 	"regexp"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/BurntSushi/toml"
 	"github.com/git-pkgs/proxy/internal/retention"
@@ -22,6 +23,11 @@ const (
 	juliaGeneralRegistryUUID = "23338594-aafe-5451-b93e-139f81909106"
 	juliaArtifactName        = "_artifact"
 	juliaRegistryName        = "_registry"
+
+	// juliaRefreshTimeout bounds a background reload of the name map.
+	juliaRefreshTimeout = 5 * time.Minute
+	// maxJuliaRegistryTomlSize bounds Registry.toml. General's is a few MB.
+	maxJuliaRegistryTomlSize = 64 << 20
 )
 
 var (
@@ -110,7 +116,7 @@ func (h *JuliaHandler) handleRegistry(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	go h.refreshNamesFromRegistry(uuid, hash)
+	go h.refreshNamesFromRegistry(context.WithoutCancel(r.Context()), uuid, hash)
 
 	ServeArtifactRequest(w, r, result)
 }
@@ -211,24 +217,16 @@ func (h *JuliaHandler) loadNamesFromUpstream(ctx context.Context) error {
 	return h.loadRegistryTarball(ctx, juliaGeneralRegistryUUID, hash)
 }
 
-// fetchGeneralRegistryHash reads /registries and returns the current tree hash
-// for the General registry.
+// fetchGeneralRegistryHash reads /registries, through the metadata cache that
+// also serves it to clients, and returns the current tree hash for the
+// General registry.
 func (h *JuliaHandler) fetchGeneralRegistryHash(ctx context.Context) (string, error) {
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, h.upstreamURL+"/registries", nil)
-	if err != nil {
-		return "", err
-	}
-	resp, err := h.proxy.HTTPClient.Do(req)
+	body, _, err := h.proxy.FetchOrCacheMetadata(ctx, "julia", "registries", h.upstreamURL+"/registries", "*/*")
 	if err != nil {
-		return "", err
-	}
-	defer func() { _ = resp.Body.Close() }()
-
-	if resp.StatusCode != http.StatusOK {
-		return "", fmt.Errorf("upstream /registries returned %d", resp.StatusCode)
+		return "", fmt.Errorf("reading /registries: %w", err)
 	}
 
-	scanner := bufio.NewScanner(resp.Body)
+	scanner := bufio.NewScanner(bytes.NewReader(body))
 	for scanner.Scan() {
 		uuid, hash, ok := parseRegistryLine(scanner.Text())
 		if ok && uuid == juliaGeneralRegistryUUID {
@@ -241,23 +239,41 @@ func (h *JuliaHandler) fetchGeneralRegistryHash(ctx context.Context) (string, er
 	return "", fmt.Errorf("general registry not listed in /registries")
 }
 
-// refreshNamesFromRegistry reloads the UUID→name map from a registry tarball
-// that has just been cached. Errors are logged but do not affect the response.
-func (h *JuliaHandler) refreshNamesFromRegistry(uuid, hash string) {
-	if uuid != juliaGeneralRegistryUUID {
+// refreshNamesFromRegistry reloads the UUID→name map after a General registry
+// tarball has been served. Only the hash /registries currently lists replaces
+// the map, so requests for older snapshots neither roll it back nor trigger a
+// download each. One reload runs at a time; a request that finds one running
+// skips its own. Errors are logged but do not affect the response.
+func (h *JuliaHandler) refreshNamesFromRegistry(ctx context.Context, uuid, hash string) {
+	if uuid != juliaGeneralRegistryUUID || h.currentNamesHash() == hash {
 		return
 	}
-	h.mu.RLock()
-	current := h.namesHash
-	h.mu.RUnlock()
-	if current == hash {
+	if !h.loadMu.TryLock() {
+		return
+	}
+	defer h.loadMu.Unlock()
+
+	ctx, cancel := context.WithTimeout(ctx, juliaRefreshTimeout)
+	defer cancel()
+	advertised, err := h.fetchGeneralRegistryHash(ctx)
+	if err != nil {
+		h.proxy.Logger.Warn("failed to refresh julia name map", "error", err)
+		return
+	}
+	if advertised != hash || h.currentNamesHash() == hash {
 		return
 	}
-	if err := h.loadRegistryTarball(context.Background(), uuid, hash); err != nil {
+	if err := h.loadRegistryTarball(ctx, uuid, hash); err != nil {
 		h.proxy.Logger.Warn("failed to refresh julia name map", "error", err)
 	}
 }
 
+func (h *JuliaHandler) currentNamesHash() string {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	return h.namesHash
+}
+
 // loadRegistryTarball downloads a registry tarball and replaces the name map
 // with the contents of its Registry.toml.
 func (h *JuliaHandler) loadRegistryTarball(ctx context.Context, uuid, hash string) error {
@@ -312,10 +328,13 @@ func extractRegistryNames(r io.Reader) (map[string]string, error) {
 			continue
 		}
 
-		data, err := io.ReadAll(tr)
+		data, err := io.ReadAll(io.LimitReader(tr, maxJuliaRegistryTomlSize+1))
 		if err != nil {
 			return nil, err
 		}
+		if len(data) > maxJuliaRegistryTomlSize {
+			return nil, fmt.Errorf("registry file Registry.toml exceeds %d bytes", maxJuliaRegistryTomlSize)
+		}
 		return parseRegistryToml(data)
 	}
 }

internal/handler/julia_refresh_test.go

@@ -0,0 +1,135 @@
+package handler
+
+import (
+	"archive/tar"
+	"bytes"
+	"compress/gzip"
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+)
+
+func juliaRegistryTarball(t *testing.T, packageName string) []byte {
+	t.Helper()
+	body := "name = \"General\"\n[packages]\n295af30f-e4ad-537b-8983-00126c2a3abe = { name = \"" + packageName + "\", path = \"R/X\" }\n"
+	var buf bytes.Buffer
+	gw := gzip.NewWriter(&buf)
+	tw := tar.NewWriter(gw)
+	if err := tw.WriteHeader(&tar.Header{Name: "Registry.toml", Mode: 0o644, Size: int64(len(body))}); err != nil {
+		t.Fatal(err)
+	}
+	_, _ = tw.Write([]byte(body))
+	_ = tw.Close()
+	_ = gw.Close()
+	return buf.Bytes()
+}
+
+type fakeJuliaUpstream struct {
+	mu         sync.Mutex
+	advertised string
+	tarballs   map[string][]byte
+	downloads  map[string]int
+}
+
+func (u *fakeJuliaUpstream) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	u.mu.Lock()
+	defer u.mu.Unlock()
+	if r.URL.Path == "/registries" {
+		_, _ = w.Write([]byte("/registry/" + juliaGeneralRegistryUUID + "/" + u.advertised + "\n"))
+		return
+	}
+	hash := r.URL.Path[strings.LastIndex(r.URL.Path, "/")+1:]
+	body, ok := u.tarballs[hash]
+	if !ok {
+		http.NotFound(w, r)
+		return
+	}
+	u.downloads[hash]++
+	_, _ = w.Write(body)
+}
+
+func (u *fakeJuliaUpstream) count(hash string) int {
+	u.mu.Lock()
+	defer u.mu.Unlock()
+	return u.downloads[hash]
+}
+
+func TestJuliaRefreshOnlyLoadsAdvertisedRegistry(t *testing.T) {
+	oldHash, newHash := strings.Repeat("a", 40), strings.Repeat("b", 40)
+	up := &fakeJuliaUpstream{
+		advertised: newHash,
+		tarballs: map[string][]byte{
+			oldHash: juliaRegistryTarball(t, "OldName"),
+			newHash: juliaRegistryTarball(t, "NewName"),
+		},
+		downloads: map[string]int{},
+	}
+	server := httptest.NewServer(up)
+	defer server.Close()
+
+	proxy, _, _, _ := setupTestProxy(t)
+	h := NewJuliaHandlerWithUpstream(proxy, server.URL)
+	ctx := context.Background()
+
+	// An older snapshot neither replaces the map nor is downloaded for it.
+	for range 5 {
+		h.refreshNamesFromRegistry(ctx, juliaGeneralRegistryUUID, oldHash)
+	}
+	if n := up.count(oldHash); n != 0 {
+		t.Fatalf("old registry downloaded %d times for the name map", n)
+	}
+
+	h.refreshNamesFromRegistry(ctx, juliaGeneralRegistryUUID, newHash)
+	if got := h.resolveName("295af30f-e4ad-537b-8983-00126c2a3abe"); got != "NewName" {
+		t.Fatalf("name = %q, want NewName", got)
+	}
+
+	// The loaded hash is not fetched again, and an old one doesn't roll back.
+	h.refreshNamesFromRegistry(ctx, juliaGeneralRegistryUUID, newHash)
+	h.refreshNamesFromRegistry(ctx, juliaGeneralRegistryUUID, oldHash)
+	if n := up.count(newHash); n != 1 {
+		t.Errorf("current registry downloaded %d times, want 1", n)
+	}
+	if got := h.resolveName("295af30f-e4ad-537b-8983-00126c2a3abe"); got != "NewName" {
+		t.Errorf("name rolled back to %q", got)
+	}
+}
+
+func TestJuliaRefreshSkipsWhileLoading(t *testing.T) {
+	proxy, _, _, _ := setupTestProxy(t)
+	h := NewJuliaHandlerWithUpstream(proxy, "http://127.0.0.1:1")
+	h.loadMu.Lock()
+	done := make(chan struct{})
+	go func() {
+		h.refreshNamesFromRegistry(context.Background(), juliaGeneralRegistryUUID, strings.Repeat("c", 40))
+		close(done)
+	}()
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("refresh waited for a running load instead of skipping")
+	}
+	h.loadMu.Unlock()
+}
+
+func TestJuliaExtractRegistryNamesLimitsRegistryToml(t *testing.T) {
+	var buf bytes.Buffer
+	gw := gzip.NewWriter(&buf)
+	tw := tar.NewWriter(gw)
+	size := int64(maxJuliaRegistryTomlSize + 1)
+	if err := tw.WriteHeader(&tar.Header{Name: "Registry.toml", Mode: 0o644, Size: size}); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := tw.Write(bytes.Repeat([]byte("#"), int(size))); err != nil {
+		t.Fatal(err)
+	}
+	_ = tw.Close()
+	_ = gw.Close()
+	if _, err := extractRegistryNames(&buf); err == nil {
+		t.Fatal("oversized Registry.toml was accepted")
+	}
+}