carlos / platform Public

Download this file

1	package lease
2
3	import (
4	"context"
5	"encoding/json"
6	"errors"
7	"fmt"
8	"io"
9	"net/http"
10	"net/http/httptest"
11	"strings"
12	"sync"
13	"testing"
14	"time"
15
16	"github.com/benbjohnson/litestream"
17	)
18
19	// casS3 is one key's worth of S3 with conditional writes: If-None-Match: *
20	// refuses an existing object, If-Match refuses another etag, with 412.
21	type casS3 struct {
22	mu sync.Mutex
23	reqs []string
24	body []byte
25	etag string
26	n int
27	puts []http.Header
28	}
29
30	func newCasS3(t *testing.T) (*casS3, string) {
31	t.Helper()
32	f := &casS3{}
33	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
34	f.mu.Lock()
35	defer f.mu.Unlock()
36	f.reqs = append(f.reqs, r.Method)
37	fail := func() {
38	w.Header().Set("Content-Type", "application/xml")
39	w.WriteHeader(http.StatusPreconditionFailed)
40	_, _ = w.Write([]byte(`<?xml version="1.0" encoding="UTF-8"?><Error><Code>PreconditionFailed</Code><Message>precondition</Message></Error>`))
41	}
42	switch r.Method {
43	case http.MethodGet:
44	if f.body == nil {
45	w.Header().Set("Content-Type", "application/xml")
46	w.WriteHeader(http.StatusNotFound)
47	_, _ = w.Write([]byte(`<?xml version="1.0" encoding="UTF-8"?><Error><Code>NoSuchKey</Code><Message>no such key</Message></Error>`))
48	return
49	}
50	w.Header().Set("ETag", f.etag)
51	_, _ = w.Write(f.body)
52	case http.MethodPut:
53	f.puts = append(f.puts, r.Header.Clone())
54	if r.Header.Get("If-None-Match") == "*" && f.body != nil {
55	fail()
56	return
57	}
58	if m := r.Header.Get("If-Match"); m != "" && strings.Trim(m, `"`) != strings.Trim(f.etag, `"`) {
59	fail()
60	return
61	}
62	b, _ := io.ReadAll(r.Body)
63	f.n++
64	f.body, f.etag = b, fmt.Sprintf(`"etag-%d"`, f.n)
65	w.Header().Set("ETag", f.etag)
66	case http.MethodDelete:
67	if m := r.Header.Get("If-Match"); m != "" && strings.Trim(m, `"`) != strings.Trim(f.etag, `"`) {
68	fail()
69	return
70	}
71	f.body, f.etag = nil, ""
72	w.WriteHeader(http.StatusNoContent)
73	}
74	}))
75	t.Cleanup(srv.Close)
76	return f, srv.URL
77	}
78
79	func (f *casS3) set(t *testing.T, l litestream.Lease) {
80	t.Helper()
81	raw, err := json.Marshal(l)
82	if err != nil {
83	t.Fatal(err)
84	}
85	f.mu.Lock()
86	defer f.mu.Unlock()
87	f.n++
88	f.body, f.etag = raw, fmt.Sprintf(`"etag-%d"`, f.n)
89	}
90
91	func (f *casS3) methods() []string {
92	f.mu.Lock()
93	defer f.mu.Unlock()
94	return append([]string(nil), f.reqs...)
95	}
96
97	func (f *casS3) lease(t *testing.T) litestream.Lease {
98	t.Helper()
99	f.mu.Lock()
100	defer f.mu.Unlock()
101	var l litestream.Lease
102	if err := json.Unmarshal(f.body, &l); err != nil {
103	t.Fatalf("lock %q: %v", f.body, err)
104	}
105	return l
106	}
107
108	func casEnv(t *testing.T) (*casS3, context.Context) {
109	t.Helper()
110	f, url := newCasS3(t)
111	fakeCreds(t)
112	t.Setenv("CARLOS_S3_ENDPOINT", url)
113	return f, context.Background()
114	}
115
116	// A released lease deletes its lock, so the next wake finds none: it
117	// writes the lease create-only in one round trip, the same lease
118	// litestream writes for an absent lock (generation 1, this owner).
119	func TestAcquireAnAbsentLockIsOneCreateOnlyWrite(t *testing.T) {
120	f, ctx := casEnv(t)
121	held, err := Acquire(ctx, "b", "inst/x.example.com", "eu-west-1")
122	if err != nil {
123	t.Fatal(err)
124	}
125	if got := f.methods(); len(got) != 1 || got[0] != http.MethodPut {
126	t.Fatalf("requests = %v, want one PUT", got)
127	}
128	if h := f.puts[0]; h.Get("If-None-Match") != "*" || h.Get("If-Match") != "" {
129	t.Fatalf("the write was not create-only: %v", h)
130	}
131	l := f.lease(t)
132	if l.Generation != 1 || l.Owner == "" || time.Until(l.ExpiresAt) < LeaseTTL-time.Minute {
133	t.Fatalf("lease = %+v, want generation 1, an owner and a full TTL", l)
134	}
135	// The handle carries the etag its write returned: release is
136	// conditional on it, so it deletes this lock and no other.
137	if err := held.Release(ctx); err != nil {
138	t.Fatal(err)
139	}
140	if got := f.methods(); got[len(got)-1] != http.MethodDelete {
141	t.Fatalf("requests = %v, want the release's DELETE", got)
142	}
143	}
144
145	// An expired lock is taken over as litestream takes it over: read, then a
146	// conditional write of the next generation.
147	func TestAcquireAnExpiredLockFallsBackToTheTakeover(t *testing.T) {
148	f, ctx := casEnv(t)
149	f.set(t, litestream.Lease{Generation: 4, ExpiresAt: time.Now().Add(-time.Minute), Owner: "gone:1"})
150	held, err := Acquire(ctx, "b", "inst/x.example.com", "eu-west-1")
151	if err != nil {
152	t.Fatal(err)
153	}
154	defer held.Release(ctx)
155	if l := f.lease(t); l.Generation != 5 || l.Owner == "gone:1" {
156	t.Fatalf("lease = %+v, want generation 5 under this owner", l)
157	}
158	}
159
160	// A live lock is someone else's: refused with its holder, as before, and
161	// left exactly as it was.
162	func TestAcquireALiveLockIsRefused(t *testing.T) {
163	f, ctx := casEnv(t)
164	f.set(t, litestream.Lease{Generation: 2, ExpiresAt: time.Now().Add(10 * time.Minute), Owner: "other:7"})
165	before := f.lease(t)
166	_, err := Acquire(ctx, "b", "inst/x.example.com", "eu-west-1")
167	var held *HeldError
168	if !errors.As(err, &held) || held.Owner != "other:7" {
169	t.Fatalf("Acquire = %v, want a HeldError naming other:7", err)
170	}
171	if after := f.lease(t); after != before {
172	t.Fatalf("lock changed: %+v -> %+v", before, after)
173	}
174	}
175