| new file mode 100644 |
| index 000000000..15dd29c20 |
| --- /dev/null |
| +++ b/internal/lease/mirror.go |
| @@ -0,0 +1,153 @@ |
| +package lease |
| + |
| +import ( |
| + "context" |
| + "fmt" |
| + "net/url" |
| + "sort" |
| + |
| + awss3 "github.com/aws/aws-sdk-go-v2/service/s3" |
| +) |
| + |
| +// The staging mirror (issue #331, instance-move spec): a residence move |
| +// copies the instance's REPLICA prefix — the litestream generations the |
| +// activator already maintains — never the database itself. MirrorPrefix is |
| +// the one delta round; the move worker calls it repeatedly until a round |
| +// reports nothing changed, and once more after the child is down. |
| +// |
| +// The interface exists for the tests (an in-memory fake) and because the |
| +// destination bucket lives in ANOTHER region: the implementation holds one |
| +// client per side. |
| + |
| +// MirrorStore is the three S3 actions a delta round needs. List answers |
| +// key→ETag under a prefix (keys RELATIVE to the bucket root, not the |
| +// prefix); Copy is a server-side CopyObject; Delete removes one key. |
| +type MirrorStore interface { |
| + List(ctx context.Context, bucket, prefix string) (map[string]string, error) |
| + Copy(ctx context.Context, srcBucket, srcKey, dstBucket, dstKey string) error |
| + Delete(ctx context.Context, bucket, key string) error |
| +} |
| + |
| +// MirrorPrefix makes dstBucket/dstPrefix an exact mirror of |
| +// srcBucket/srcPrefix: copy keys that are new or whose ETag differs, delete |
| +// keys the source no longer has (litestream retention must hold in staging |
| +// too — a restore that replays pruned WAL segments is a restore nobody |
| +// tested). Returns how many objects changed; 0 means the sides were |
| +// identical when listed — the move worker's "caught up". |
| +// |
| +// ETag equality is a valid sameness test here even for multipart uploads: |
| +// both sides of a server-side COPY carry the copy's own single-part ETag, |
| +// and litestream's own objects are single PUTs. |
| +func MirrorPrefix(ctx context.Context, m MirrorStore, srcBucket, srcPrefix, dstBucket, dstPrefix string) (int, error) { |
| + src, err := m.List(ctx, srcBucket, srcPrefix) |
| + if err != nil { |
| + return 0, fmt.Errorf("mirror: list source %s/%s: %w", srcBucket, srcPrefix, err) |
| + } |
| + dst, err := m.List(ctx, dstBucket, dstPrefix) |
| + if err != nil { |
| + return 0, fmt.Errorf("mirror: list staging %s/%s: %w", dstBucket, dstPrefix, err) |
| + } |
| + // Deterministic order: a failed round resumes the same way everywhere. |
| + rel := make([]string, 0, len(src)) |
| + for k := range src { |
| + rel = append(rel, k[len(srcPrefix):]) |
| + } |
| + sort.Strings(rel) |
| + changed := 0 |
| + for _, r := range rel { |
| + sk, dk := srcPrefix+r, dstPrefix+r |
| + if et, ok := dst[dk]; ok && et == src[sk] { |
| + continue |
| + } |
| + if err := m.Copy(ctx, srcBucket, sk, dstBucket, dk); err != nil { |
| + return changed, fmt.Errorf("mirror: copy %s: %w", sk, err) |
| + } |
| + changed++ |
| + } |
| + for dk := range dst { |
| + r := dk[len(dstPrefix):] |
| + if _, ok := src[srcPrefix+r]; !ok { |
| + if err := m.Delete(ctx, dstBucket, dk); err != nil { |
| + return changed, fmt.Errorf("mirror: delete %s: %w", dk, err) |
| + } |
| + changed++ |
| + } |
| + } |
| + return changed, nil |
| +} |
| + |
| +// s3Mirror implements MirrorStore over two regional clients: reads and the |
| +// copy call go to the DESTINATION region's client (CopyObject executes in |
| +// the destination bucket's region and pulls from the source), source |
| +// listings to the source's. |
| +type s3Mirror struct { |
| + src, dst *awss3.Client |
| +} |
| + |
| +// NewS3Mirror builds the two-region pair with the same credential chain and |
| +// endpoint compatibility every other S3 consumer here uses (endpoint.go's |
| +// newLeaserClient). |
| +func NewS3Mirror(ctx context.Context, srcRegion, dstRegion string) (MirrorStore, error) { |
| + src, _, err := newLeaserClient(ctx, srcRegion) |
| + if err != nil { |
| + return nil, err |
| + } |
| + dst, _, err := newLeaserClient(ctx, dstRegion) |
| + if err != nil { |
| + return nil, err |
| + } |
| + return &s3Mirror{src: src, dst: dst}, nil |
| +} |
| + |
| +func (m *s3Mirror) client(bucketSide *awss3.Client) *awss3.Client { return bucketSide } |
| + |
| +func (m *s3Mirror) list(ctx context.Context, c *awss3.Client, bucket, prefix string) (map[string]string, error) { |
| + out := map[string]string{} |
| + var token *string |
| + for { |
| + resp, err := c.ListObjectsV2(ctx, &awss3.ListObjectsV2Input{ |
| + Bucket: &bucket, Prefix: &prefix, ContinuationToken: token, |
| + }) |
| + if err != nil { |
| + return nil, err |
| + } |
| + for _, o := range resp.Contents { |
| + if o.Key == nil { |
| + continue |
| + } |
| + et := "" |
| + if o.ETag != nil { |
| + et = *o.ETag |
| + } |
| + out[*o.Key] = et |
| + } |
| + if resp.IsTruncated == nil || !*resp.IsTruncated || resp.NextContinuationToken == nil { |
| + return out, nil |
| + } |
| + token = resp.NextContinuationToken |
| + } |
| +} |
| + |
| +// List tries the destination client first and falls back to the source's — |
| +// the two buckets live in different regions and the caller does not say |
| +// which side a bucket is on. A region-mismatch answer from S3 is an error, |
| +// so the fallback is unambiguous. |
| +func (m *s3Mirror) List(ctx context.Context, bucket, prefix string) (map[string]string, error) { |
| + if out, err := m.list(ctx, m.src, bucket, prefix); err == nil { |
| + return out, nil |
| + } |
| + return m.list(ctx, m.dst, bucket, prefix) |
| +} |
| + |
| +func (m *s3Mirror) Copy(ctx context.Context, srcBucket, srcKey, dstBucket, dstKey string) error { |
| + source := url.PathEscape(srcBucket + "/" + srcKey) |
| + _, err := m.dst.CopyObject(ctx, &awss3.CopyObjectInput{ |
| + Bucket: &dstBucket, Key: &dstKey, CopySource: &source, |
| + }) |
| + return err |
| +} |
| + |
| +func (m *s3Mirror) Delete(ctx context.Context, bucket, key string) error { |
| + _, err := m.dst.DeleteObject(ctx, &awss3.DeleteObjectInput{Bucket: &bucket, Key: &key}) |
| + return err |
| +} |
amadan