amadan
relay: heartbeat overwrites by read-then-CAS (review fix)
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
2 files changed,
+52
−2
internal/relay/authority.go+28 −2internal/relay/authority_test.go+24 −0
diff --git a/internal/relay/authority.go b/internal/relay/authority.go| index 0c96456cc..74937be56 100644 |
| --- a/internal/relay/authority.go |
| +++ b/internal/relay/authority.go |
| @@ -73,11 +73,37 @@ type Heartbeat struct { |
| MastersError string `json:"masters_error,omitempty"` |
| } |
| +// WriteHeartbeat writes a box's heartbeat by read-then-CAS: Get the current |
| +// etag (absent means create with ifMatch ""), then Put with it. A box owns |
| +// its own heartbeat key, so an ErrCASConflict can only be our own previous |
| +// write racing a restart — retry the whole sequence once more and return |
| +// whatever the second attempt reports. |
| func WriteHeartbeat(ctx context.Context, s release.Store, hb Heartbeat) error { |
| body, err := json.Marshal(hb) |
| if err != nil { |
| return err |
| } |
| - _, err = s.Put(ctx, HeartbeatKey(hb.Box), body, "") |
| - return err |
| + key := HeartbeatKey(hb.Box) |
| + var lastErr error |
| + for attempt := 0; attempt < 2; attempt++ { |
| + etag := "" |
| + _, curEtag, err := s.Get(ctx, key) |
| + switch { |
| + case errors.Is(err, os.ErrNotExist): |
| + // no existing record: create-only Put below. |
| + case err != nil: |
| + return err |
| + default: |
| + etag = curEtag |
| + } |
| + if _, err := s.Put(ctx, key, body, etag); err != nil { |
| + lastErr = err |
| + if errors.Is(err, release.ErrCASConflict) { |
| + continue |
| + } |
| + return err |
| + } |
| + return nil |
| + } |
| + return lastErr |
| } |
diff --git a/internal/relay/authority_test.go b/internal/relay/authority_test.go| index 34310f48a..a06be31ef 100644 |
| --- a/internal/relay/authority_test.go |
| +++ b/internal/relay/authority_test.go |
| @@ -30,6 +30,30 @@ func TestAuthorityRoundTrip(t *testing.T) { |
| } |
| } |
| +func TestWriteHeartbeatTwice(t *testing.T) { |
| + s := newMemStore(t) |
| + ctx := context.Background() |
| + t0 := time.Unix(1_800_000_000, 0).UTC() |
| + if err := WriteHeartbeat(ctx, s, Heartbeat{Box: "box1", Region: "eu-west-1", ReadAt: t0}); err != nil { |
| + t.Fatalf("first write: %v", err) |
| + } |
| + t1 := t0.Add(30 * time.Second) |
| + if err := WriteHeartbeat(ctx, s, Heartbeat{Box: "box1", Region: "eu-west-1", ReadAt: t1}); err != nil { |
| + t.Fatalf("second write: %v", err) |
| + } |
| + body, _, err := s.Get(ctx, HeartbeatKey("box1")) |
| + if err != nil { |
| + t.Fatal(err) |
| + } |
| + var got Heartbeat |
| + if err := json.Unmarshal(body, &got); err != nil { |
| + t.Fatal(err) |
| + } |
| + if !got.ReadAt.Equal(t1) { |
| + t.Fatalf("read_at = %v, want %v", got.ReadAt, t1) |
| + } |
| +} |
| + |
| func TestLease(t *testing.T) { |
| var l Lease |
| t0 := time.Unix(1_800_000_000, 0) |