Serialise gem push and yank per version with a Redis lock

internal/api/v1/gems.go

@@ -30,7 +30,7 @@ import (
 
 var (
 	gemsMaxPushBytes int64 = 256 << 20
-	gemsPushHook           = func() {}
+	gemsLockWait           = 30 * time.Second
 )
 
 const (
@@ -53,6 +53,16 @@ func (h *ProxyHandler) localGemsRemote(w http.ResponseWriter, r *http.Request) (
 	return remote, true
 }
 
+// lockGemVersion serialises push and yank of one gem version across replicas.
+func (h *ProxyHandler) lockGemVersion(w http.ResponseWriter, r *http.Request, repo, full string) (func(), bool) {
+	unlock, err := h.engine.Lock(r.Context(), repo, "_gemlock/"+full, gemsLockWait, 5*time.Minute)
+	if err != nil {
+		gemsText(w, http.StatusServiceUnavailable, full+" is being pushed or yanked, retry")
+		return nil, false
+	}
+	return unlock, true
+}
+
 func (h *ProxyHandler) handleLocalPost(w http.ResponseWriter, r *http.Request) {
 	if chi.URLParam(r, "*") != "api/v1/gems" {
 		gemsText(w, http.StatusNotFound, "not found")
@@ -79,60 +89,61 @@ func (h *ProxyHandler) handleLocalPost(w http.ResponseWriter, r *http.Request) {
 	}
 
 	full := spec.FullName()
+	unlock, ok := h.lockGemVersion(w, r, remote.Name, full)
+	if !ok {
+		return
+	}
+	defer unlock()
 	conflict := func() {
 		gemsText(w, http.StatusConflict, fmt.Sprintf("%s already exists; overwrites are not allowed", full))
 	}
 	index := gemsIndexPrefix + spec.Name + "/" + full
-	paths := []string{index, "gems/" + full + ".gem", "quick/Marshal.4.8/" + full + ".gemspec.rz"}
-	yanked := func() bool {
-		f, err := h.db.GetLocalFile(r.Context(), remote.Name, gemsYankedPrefix+spec.Name+"/"+full)
+	for _, p := range []string{gemsYankedPrefix + spec.Name + "/" + full, index} {
+		f, err := h.db.GetLocalFile(r.Context(), remote.Name, p)
 		if err != nil {
 			gemsText(w, http.StatusInternalServerError, err.Error())
-			return true
+			return
+		}
+		if f != nil && p == index {
+			conflict()
+			return
 		}
 		if f != nil {
 			gemsText(w, http.StatusConflict, fmt.Sprintf("%s was yanked; re-pushing a yanked version is not allowed", full))
+			return
 		}
-		return f != nil
 	}
-	if yanked() {
-		return
+	// The version is not live, so on failure every path is an invisible leftover and safe to drop.
+	paths := []string{"gems/" + full + ".gem", "quick/Marshal.4.8/" + full + ".gemspec.rz", index}
+	fail := func(status int, msg string) {
+		for _, p := range paths {
+			if err := h.db.DeleteLocalFile(context.WithoutCancel(r.Context()), remote.Name, p); err != nil {
+				slog.Error("gem push cleanup failed", "repo", remote.Name, "path", p, "error", err)
+			}
+		}
+		gemsText(w, status, msg)
 	}
-	gemsPushHook()
 	// The index line gates visibility, so it is written last.
 	for _, f := range []struct {
 		path string
 		data []byte
 	}{
-		{paths[1], gem},
-		{paths[2], quick},
+		{paths[0], gem},
+		{paths[1], quick},
 	} {
 		ok, err := h.npmStoreImmutable(r.Context(), remote.Name, f.path, f.data, "application/octet-stream")
 		if err != nil {
-			gemsText(w, http.StatusInternalServerError, err.Error())
+			fail(http.StatusInternalServerError, err.Error())
 			return
 		}
 		if !ok {
-			conflict()
+			fail(http.StatusConflict, fmt.Sprintf("%s already exists; overwrites are not allowed", full))
 			return
 		}
 	}
 	line := gems.InfoLine(spec, gem)
 	if _, err := h.npmStore(r.Context(), remote.Name, index, []byte(line), "text/plain", false); err != nil {
-		if errors.Is(err, database.ErrAlreadyExists) {
-			conflict()
-			return
-		}
-		gemsText(w, http.StatusInternalServerError, err.Error())
-		return
-	}
-	// A yank that raced this push tombstoned the version before or after its deletes ran; undo our writes.
-	if yanked() {
-		for _, p := range paths {
-			if err := h.db.DeleteLocalFile(r.Context(), remote.Name, p); err != nil {
-				slog.Error("gem push rollback failed", "repo", remote.Name, "path", p, "error", err)
-			}
-		}
+		fail(http.StatusInternalServerError, err.Error())
 		return
 	}
 
@@ -174,6 +185,11 @@ func (h *ProxyHandler) handleLocalDelete(w http.ResponseWriter, r *http.Request)
 	}
 
 	full := spec.FullName()
+	unlock, ok := h.lockGemVersion(w, r, remote.Name, full)
+	if !ok {
+		return
+	}
+	defer unlock()
 	index := gemsIndexPrefix + spec.Name + "/" + full
 	existing, err := h.db.GetLocalFile(r.Context(), remote.Name, index)
 	if err != nil {
@@ -195,7 +211,7 @@ func (h *ProxyHandler) handleLocalDelete(w http.ResponseWriter, r *http.Request)
 		return
 	}
 	for _, p := range []string{index, "gems/" + full + ".gem", "quick/Marshal.4.8/" + full + ".gemspec.rz"} {
-		if err := h.db.DeleteLocalFile(r.Context(), remote.Name, p); err != nil {
+		if err := h.db.DeleteLocalFile(context.WithoutCancel(r.Context()), remote.Name, p); err != nil {
 			gemsText(w, http.StatusInternalServerError, err.Error())
 			return
 		}

internal/api/v1/gems_handler_test.go

@@ -4,14 +4,18 @@ import (
 	"archive/tar"
 	"bytes"
 	"compress/gzip"
+	"context"
 	"crypto/md5"
 	"crypto/sha256"
 	"encoding/hex"
 	"fmt"
 	"net/http"
 	"strings"
+	"sync"
 	"testing"
+	"time"
 
+	"git.unkin.net/unkin/artifactapi/internal/cache"
 	"git.unkin.net/unkin/artifactapi/pkg/models"
 )
 
@@ -176,28 +180,74 @@ func TestGemsPushRejects(t *testing.T) {
 	}
 }
 
-func TestGemsPushRacingYank(t *testing.T) {
+func TestGemsPushYankSerialised(t *testing.T) {
 	h, srv := npmTestServer(t)
 	repo := npmRepo(t, h, models.PackageGems)
 	base := srv.URL + "/api/v1/local/" + repo + "/"
-	g := testGemFile(t, "race", "1.0.0", "ruby")
-	if code, body := npmDo(t, http.MethodPost, base+"api/v1/gems", string(g)); code != http.StatusOK {
-		t.Fatalf("push = %d: %s", code, body)
+	g := string(testGemFile(t, "race", "1.0.0", "ruby"))
+	var wg sync.WaitGroup
+	for i := 0; i < 16; i++ {
+		wg.Add(2)
+		go func() {
+			defer wg.Done()
+			if code, body := npmDo(t, http.MethodPost, base+"api/v1/gems", g); code != http.StatusOK && code != http.StatusConflict {
+				t.Errorf("push = %d: %s", code, body)
+			}
+		}()
+		go func() {
+			defer wg.Done()
+			if code, body := npmDo(t, http.MethodDelete, base+"api/v1/gems/yank?gem_name=race&version=1.0.0", ""); code != http.StatusOK && code != http.StatusNotFound {
+				t.Errorf("yank = %d: %s", code, body)
+			}
+		}()
 	}
+	wg.Wait()
 
-	gemsPushHook = func() {
-		gemsPushHook = func() {}
-		if code, body := npmDo(t, http.MethodDelete, base+"api/v1/gems/yank?gem_name=race&version=1.0.0", ""); code != http.StatusOK {
-			t.Errorf("yank = %d: %s", code, body)
+	present := func(p string) bool {
+		f, err := h.db.GetLocalFile(context.Background(), repo, p)
+		if err != nil {
+			t.Fatal(err)
+		}
+		return f != nil
+	}
+	live := []string{"_index/race/race-1.0.0", "gems/race-1.0.0.gem", "quick/Marshal.4.8/race-1.0.0.gemspec.rz"}
+	tomb := present("_yanked/race/race-1.0.0")
+	for _, p := range live {
+		if present(p) == tomb {
+			t.Errorf("tombstoned=%v but %s present=%v", tomb, p, !tomb)
 		}
 	}
-	t.Cleanup(func() { gemsPushHook = func() {} })
-	if code, body := npmDo(t, http.MethodPost, base+"api/v1/gems", string(g)); code != http.StatusConflict || !strings.Contains(body, "yanked") {
-		t.Fatalf("push racing yank = %d %q, want 409", code, body)
-	}
-	for _, p := range []string{"info/race", "gems/race-1.0.0.gem", "quick/Marshal.4.8/race-1.0.0.gemspec.rz", "_index/race/race-1.0.0", "_yanked/race/race-1.0.0"} {
-		if code, _ := npmDo(t, http.MethodGet, base+p, ""); code != http.StatusNotFound {
-			t.Errorf("GET %s = %d, want 404", p, code)
-		}
+	_, info := npmDo(t, http.MethodGet, base+"info/race", "")
+	if strings.Contains(info, "1.0.0") == tomb {
+		t.Errorf("tombstoned=%v but info = %q", tomb, info)
+	}
+}
+
+func TestGemsLockTimeoutIs503(t *testing.T) {
+	h, srv := npmTestServer(t)
+	repo := npmRepo(t, h, models.PackageGems)
+	base := srv.URL + "/api/v1/local/" + repo + "/"
+	redis, err := cache.NewRedis(testRedis)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if ok, err := redis.AcquireLock(context.Background(), repo, "_gemlock/busy-1.0.0", time.Minute); !ok || err != nil {
+		t.Fatalf("hold lock = %v %v", ok, err)
+	}
+	prev := gemsLockWait
+	gemsLockWait = 100 * time.Millisecond
+	t.Cleanup(func() { gemsLockWait = prev })
+
+	if code, _ := npmDo(t, http.MethodPost, base+"api/v1/gems", string(testGemFile(t, "busy", "1.0.0", "ruby"))); code != http.StatusServiceUnavailable {
+		t.Fatalf("push under held lock = %d, want 503", code)
+	}
+	if code, _ := npmDo(t, http.MethodDelete, base+"api/v1/gems/yank?gem_name=busy&version=1.0.0", ""); code != http.StatusServiceUnavailable {
+		t.Fatalf("yank under held lock = %d, want 503", code)
+	}
+	if err := redis.ReleaseLock(context.Background(), repo, "_gemlock/busy-1.0.0"); err != nil {
+		t.Fatal(err)
+	}
+	if code, body := npmDo(t, http.MethodPost, base+"api/v1/gems", string(testGemFile(t, "busy", "1.0.0", "ruby"))); code != http.StatusOK {
+		t.Fatalf("push after release = %d: %s", code, body)
 	}
 }

internal/api/v1/npm_handler_test.go

@@ -15,7 +15,9 @@ import (
 
 	"github.com/go-chi/chi/v5"
 
+	"git.unkin.net/unkin/artifactapi/internal/cache"
 	"git.unkin.net/unkin/artifactapi/internal/database"
+	"git.unkin.net/unkin/artifactapi/internal/proxy"
 	"git.unkin.net/unkin/artifactapi/internal/storage"
 	"git.unkin.net/unkin/artifactapi/internal/testsupport"
 	"git.unkin.net/unkin/artifactapi/internal/virtual"
@@ -25,6 +27,7 @@ import (
 var (
 	testDSN   string
 	testMinio testsupport.MinioConn
+	testRedis string
 )
 
 func TestMain(m *testing.M) {
@@ -38,8 +41,15 @@ func TestMain(m *testing.M) {
 		termPG()
 		os.Exit(m.Run())
 	}
-	testDSN, testMinio = dsn, conn
+	redisURL, termRedis, err := testsupport.StartRedis(ctx)
+	if err != nil {
+		termMinio()
+		termPG()
+		os.Exit(m.Run())
+	}
+	testDSN, testMinio, testRedis = dsn, conn, redisURL
 	code := m.Run()
+	termRedis()
 	termMinio()
 	termPG()
 	os.Exit(code)
@@ -65,7 +75,11 @@ func npmTestServer(t *testing.T) (*ProxyHandler, *httptest.Server) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	h := NewProxyHandler(nil, nil, db, store, nil)
+	redis, err := cache.NewRedis(testRedis)
+	if err != nil {
+		t.Fatal(err)
+	}
+	h := NewProxyHandler(proxy.NewEngine(db, redis, store), nil, db, store, nil)
 	r := chi.NewRouter()
 	r.Mount("/api/v1", h.Routes())
 	srv := httptest.NewServer(r)

internal/proxy/engine.go

@@ -228,10 +228,11 @@ func (e *Engine) Evict(ctx context.Context, remoteName, path string) error {
 		return fmt.Errorf("get remote: %w", err)
 	}
 	if !wildcard {
-		if err := e.waitForLock(ctx, remoteName, path); err != nil {
+		release, err := e.Lock(ctx, remoteName, path, e.evictLockWait, fetchLockTTL)
+		if err != nil {
 			return err
 		}
-		defer func() { _ = e.cache.ReleaseLock(context.WithoutCancel(ctx), remoteName, path) }()
+		defer release()
 		if err := e.db.DeleteArtifact(ctx, remoteName, path); err != nil {
 			return fmt.Errorf("delete artifact: %w", err)
 		}
@@ -252,13 +253,20 @@ func (e *Engine) Evict(ctx context.Context, remoteName, path string) error {
 	return e.cache.ForgetPrefix(ctx, remoteName, prefix)
 }
 
-// waitForLock takes the per-path fetch lock so an in-flight Fetch cannot
-// re-set TTL/ETag keys after an evict. It fails with a 503 when the lock
-// cannot be taken within evictLockWait or Redis errors.
-func (e *Engine) waitForLock(ctx context.Context, remoteName, path string) error {
-	deadline := time.Now().Add(e.evictLockWait)
+// Lock takes the per-path Redis lock (the one Fetch holds) and returns its
+// release. It fails with a 503 when the lock cannot be taken within wait or
+// Redis errors.
+func (e *Engine) Lock(ctx context.Context, remoteName, path string, wait, ttl time.Duration) (func(), error) {
+	if err := e.waitForLock(ctx, remoteName, path, wait, ttl); err != nil {
+		return nil, err
+	}
+	return func() { _ = e.cache.ReleaseLock(context.WithoutCancel(ctx), remoteName, path) }, nil
+}
+
+func (e *Engine) waitForLock(ctx context.Context, remoteName, path string, wait, ttl time.Duration) error {
+	deadline := time.Now().Add(wait)
 	for {
-		ok, err := e.cache.AcquireLock(ctx, remoteName, path, fetchLockTTL)
+		ok, err := e.cache.AcquireLock(ctx, remoteName, path, ttl)
 		if ok {
 			return nil
 		}