fix(persist): keep the shared id counter across restore

internal/idgen/counter_test.go

@@ -0,0 +1,81 @@
+package idgen
+
+import (
+	"strings"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAdvanceToMovesCounterForward(t *testing.T) {
+	Reset()
+
+	AdvanceTo(0x2a)
+	assert.Equal(t, uint64(0x2a), Counter())
+	assert.Equal(t, "i-0000002b", GenerateID("i-"))
+	assert.True(t, strings.HasSuffix(OCID("vcn", "", "us-ashburn-1"), "aaaaaaaa000000000000002c"))
+}
+
+func TestAdvanceToNeverRewinds(t *testing.T) {
+	Reset()
+	AdvanceTo(100)
+
+	AdvanceTo(5)
+	assert.Equal(t, uint64(100), Counter())
+
+	AdvanceTo(100)
+	assert.Equal(t, uint64(100), Counter())
+}
+
+// TestAdvanceToConcurrentWithGenerateID runs AdvanceTo against GenerateID from
+// many goroutines (run with -race): ids stay unique and the counter ends at
+// least at the highest target.
+func TestAdvanceToConcurrentWithGenerateID(t *testing.T) {
+	Reset()
+
+	const (
+		workers = 8
+		perW    = 500
+	)
+
+	var (
+		wg   sync.WaitGroup
+		mu   sync.Mutex
+		seen = make(map[string]bool, workers*perW)
+	)
+
+	for w := range workers {
+		wg.Add(2)
+
+		go func() {
+			defer wg.Done()
+
+			local := make([]string, 0, perW)
+			for range perW {
+				local = append(local, GenerateID("x-"))
+			}
+
+			mu.Lock()
+			defer mu.Unlock()
+
+			for _, id := range local {
+				assert.False(t, seen[id], "duplicate id %s", id)
+				seen[id] = true
+			}
+		}()
+
+		go func() {
+			defer wg.Done()
+
+			for i := range perW {
+				AdvanceTo(uint64(w*perW + i))
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	assert.Len(t, seen, workers*perW)
+	assert.GreaterOrEqual(t, Counter(), uint64(workers*perW-1))
+}

internal/idgen/idgen.go

@@ -102,6 +102,26 @@ func next() uint64 {
 	return atomic.AddUint64(&counter, 1)
 }
 
+// Counter returns the current value of the shared id counter, the suffix of
+// the most recently minted GenerateID or OCID. A snapshot records it so a
+// restore can move the counter past every id it brings back.
+func Counter() uint64 {
+	return atomic.LoadUint64(&counter)
+}
+
+// AdvanceTo moves the shared id counter up to n when it is lower, so the next
+// GenerateID or OCID is above n. It never moves the counter back: loading an
+// older snapshot into a process that has already minted more ids keeps the
+// higher value. Safe for concurrent use with GenerateID.
+func AdvanceTo(n uint64) {
+	for {
+		cur := atomic.LoadUint64(&counter)
+		if cur >= n || atomic.CompareAndSwapUint64(&counter, cur, n) {
+			return
+		}
+	}
+}
+
 // GenerateID generates an ID with the given prefix (e.g., "i-", "vpc-", "sg-").
 func GenerateID(prefix string) string {
 	return fmt.Sprintf("%s%08x", prefix, next())

persist/idcounter.go

@@ -0,0 +1,57 @@
+package persist
+
+import (
+	"regexp"
+	"strconv"
+)
+
+// legacyIDCap bounds the values the legacy scan accepts. A process never mints
+// anywhere near 2^28 counter ids, so a larger value is random hex (a UUID or
+// hash segment), not a counter id. Keeping the counter below 2^32 also keeps
+// GenerateID's suffix at the 8 hex characters SDKs expect.
+const legacyIDCap = 1 << 28
+
+// counterHexLen is the width of GenerateID's hex suffix, and of the low 32 bits
+// of an OCID suffix.
+const counterHexLen = 8
+
+// hexRun matches a run of hex digits long enough to hold a counter id.
+var hexRun = regexp.MustCompile(`[0-9a-fA-F]{8,}`)
+
+// idFloor returns the value the shared id counter must reach before ps is
+// restored: the recorded counter, or for a snapshot written before it was
+// recorded, a best-effort estimate from the ids in the state.
+func idFloor(ps *ProviderState) uint64 {
+	if ps.IDCounter != 0 || len(ps.Services) == 0 {
+		return ps.IDCounter
+	}
+
+	var floor uint64
+	for _, raw := range ps.Services {
+		floor = max(floor, legacyIDFloor(raw))
+	}
+
+	return floor
+}
+
+// legacyIDFloor estimates the highest counter value used by the ids in data.
+// It is best-effort: it takes the last 8 hex characters of every hex run of 8
+// or more (a GenerateID suffix, the tail of an OCID, the newest of several
+// concatenated ids) and ignores anything at or above legacyIDCap. Random hex
+// below the cap can only push the counter higher than needed, which is safe:
+// ids minted afterwards look larger but keep their 8 hex characters. Ids held
+// in base64-encoded []byte fields are not decoded, so they are not seen.
+func legacyIDFloor(data []byte) uint64 {
+	var floor uint64
+
+	for _, run := range hexRun.FindAll(data, -1) {
+		v, err := strconv.ParseUint(string(run[len(run)-counterHexLen:]), 16, 64)
+		if err != nil || v >= legacyIDCap {
+			continue
+		}
+
+		floor = max(floor, v)
+	}
+
+	return floor
+}

persist/idcounter_test.go

@@ -0,0 +1,53 @@
+package persist
+
+import (
+	"encoding/json"
+	"testing"
+)
+
+func TestLegacyIDFloor(t *testing.T) {
+	cases := []struct {
+		name string
+		data string
+		want uint64
+	}{
+		{"prefixed id", `{"id":"i-0000002a"}`, 0x2a},
+		{"prefix ending in a hex letter", `{"id":"ANPA0000003b"}`, 0x3b},
+		{"upper-case id", `{"id":"Z0000004C"}`, 0x4c},
+		{"concatenated ids take the newest", `{"id":"0000000100000009"}`, 9},
+		{"ocid suffix", `{"id":"ocid1.vcn.oc1.iad.aaaaaaaa0000000000000011"}`, 0x11},
+		{"id inside an arn", `{"arn":"arn:aws:acm:us-east-1:123456789012:certificate/00000020"}`, 0x20},
+		{"id inside a hostname", `{"e":"vpc-d-00000030.us-east-1.es.amazonaws.com"}`, 0x30},
+		{"map key", `{"i-00000040":{}}`, 0x40},
+		{"random hex above the cap is ignored", `{"id":"i-0000002a","u":"9f3c2b1e"}`, 0x2a},
+		{"short runs are ignored", `{"id":"abc123"}`, 0},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := legacyIDFloor([]byte(tc.data)); got != tc.want {
+				t.Fatalf("legacyIDFloor(%s) = %#x, want %#x", tc.data, got, tc.want)
+			}
+		})
+	}
+}
+
+func TestIDFloorPrefersRecordedCounter(t *testing.T) {
+	ps := &ProviderState{
+		Services:  map[string]json.RawMessage{"ec2": json.RawMessage(`{"id":"i-000000ff"}`)},
+		IDCounter: 7,
+	}
+
+	if got := idFloor(ps); got != 7 {
+		t.Fatalf("idFloor = %d, want the recorded 7", got)
+	}
+
+	ps.IDCounter = 0
+	if got := idFloor(ps); got != 0xff {
+		t.Fatalf("legacy idFloor = %#x, want 0xff", got)
+	}
+
+	if got := idFloor(&ProviderState{}); got != 0 {
+		t.Fatalf("empty idFloor = %d, want 0", got)
+	}
+}