[raftstore] Rename context methods

cmds/core-service/main.go

@@ -368,9 +368,9 @@ func RunHTTPServer(ctx context.Context, ctxCanceler func(), address, locality st
 	handler = authorizer.TokenMiddleware(handler)
 	handler = http.TimeoutHandler(handler, *timeout, "request timeout")
 	handler = logging.HTTPMiddleware(logger, *dumpRequests, handler)
-	handler = timestamp.RequestTimestampMiddleware(handler)
+	handler = timestamp.Middleware(handler)
 	handler = random.Middleware(handler)
-	handler = requestlocality.LocalityMiddleware(locality)(handler)
+	handler = requestlocality.Middleware(locality)(handler)
 
 	if *enableMetrics || *enableTracing {
 		// We use the default settings; the APIRouter handler will override the span value accordingly, as it has more information.

pkg/aux_/pool_participants.go

@@ -96,7 +96,7 @@ func (a *Server) PutDSSInstancesHeartbeat(ctx context.Context, req *restapi.PutD
 		}
 		heartbeat.Timestamp = &ts
 	} else {
-		now := timestamp.MustGetRequestTimestamp(ctx)
+		now := timestamp.MustFromContext(ctx)
 		heartbeat.Timestamp = &now
 	}
 

pkg/aux_/store/memstore/dss.go

@@ -11,7 +11,7 @@ import (
 )
 
 func (r *repo) SaveOwnMetadata(ctx context.Context, loc string, publicEndpoint string) error {
-	now := timestamp.MustGetRequestTimestamp(ctx)
+	now := timestamp.MustFromContext(ctx)
 
 	r.state.Participants[locality(loc)] = &participant{
 		PublicEndpoint: publicEndpoint,

pkg/aux_/store/memstore/dss_test.go

@@ -15,7 +15,7 @@ var fakeClock = clockwork.NewFakeClock()
 
 func TestSaveOwnMetadataRoundTrip(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	r := newRepo()
 
 	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
@@ -35,7 +35,7 @@ func TestSaveOwnMetadataRoundTrip(t *testing.T) {
 
 func TestSaveOwnMetadataUpsert(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	r := newRepo()
 
 	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://old.example.com"))
@@ -50,7 +50,7 @@ func TestSaveOwnMetadataUpsert(t *testing.T) {
 
 func TestGetDSSMetadataPicksLatestHeartbeat(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	r := newRepo()
 
 	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
@@ -71,7 +71,7 @@ func TestGetDSSMetadataPicksLatestHeartbeat(t *testing.T) {
 
 func TestGetDSSMetadataUpdatesHeartbeatPerSource(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	r := newRepo()
 
 	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))

pkg/aux_/store/memstore/snapshot_test.go

@@ -17,7 +17,7 @@ import (
 
 func TestSnapshotRoundTrip(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	src := newRepo()
 	require.NoError(t, src.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
 	ts := time.Now().UTC()
@@ -40,7 +40,7 @@ func TestSnapshotRoundTrip(t *testing.T) {
 
 func TestRestoreFromSnapshotReplacesState(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	src := newRepo()
 	require.NoError(t, src.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
 	data, err := src.GetSnapshot()

pkg/aux_/store/memstore/store_test.go

@@ -10,7 +10,7 @@ import (
 
 func TestCheckpointRestore(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 
 	r := newRepo()
 
@@ -34,7 +34,7 @@ func TestCheckpointRestore(t *testing.T) {
 
 func TestCheckpointIsolatesUpsert(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	r := newRepo()
 
 	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://old.example.com"))

pkg/locality/locality.go

@@ -7,30 +7,30 @@ import (
 	"github.com/interuss/stacktrace"
 )
 
-type localityKey struct{}
+type key struct{}
 
-// MustGetRequestLocality returns the request locality from the context and panics if it is not
+// MustFromContext returns the request locality from the context and panics if it is not
 // present, which is a programming error.
-func MustGetRequestLocality(ctx context.Context) string {
-	locality, ok := ctx.Value(localityKey{}).(string)
+func MustFromContext(ctx context.Context) string {
+	locality, ok := ctx.Value(key{}).(string)
 	if !ok {
 		panic(stacktrace.NewError("request locality not present in context"))
 	}
 
 	return locality
 }
 
-// WithRequestLocality returns a new context with the given locality.
-func WithRequestLocality(ctx context.Context, locality string) context.Context {
-	return context.WithValue(ctx, localityKey{}, locality)
+// NewContext returns a new context with the given locality.
+func NewContext(ctx context.Context, locality string) context.Context {
+	return context.WithValue(ctx, key{}, locality)
 }
 
-// LocalityMiddleware is an HTTP middleware that stamps each incoming request with this
+// Middleware is an HTTP middleware that stamps each incoming request with this
 // DSS instance's locality so that locality-dependent operations execute deterministically across nodes.
-func LocalityMiddleware(locality string) func(http.Handler) http.Handler {
+func Middleware(locality string) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			next.ServeHTTP(w, r.WithContext(WithRequestLocality(r.Context(), locality)))
+			next.ServeHTTP(w, r.WithContext(NewContext(r.Context(), locality)))
 		})
 	}
 }

pkg/raftstore/consensus/proposal.go

@@ -35,7 +35,7 @@ type Proposal struct {
 }
 
 func (c *Consensus) newProposal(ctx context.Context, requestType RequestType, value []byte, readOnly bool) Proposal {
-	timestamp := timestamp.MustGetRequestTimestamp(ctx)
+	timestamp := timestamp.MustFromContext(ctx)
 	seed := random.MustFromContext(ctx)
 
 	return Proposal{

pkg/raftstore/store.go

@@ -116,8 +116,8 @@ func (s *Store[R]) processCommits(ctx context.Context, commitCh <-chan consensus
 				continue
 			}
 
-			proposalCtx := timestamp.WithRequestTimestamp(ctx, commit.Prop.Timestamp)
-			proposalCtx = locality.WithRequestLocality(proposalCtx, commit.Prop.Locality)
+			proposalCtx := timestamp.NewContext(ctx, commit.Prop.Timestamp)
+			proposalCtx = locality.NewContext(proposalCtx, commit.Prop.Locality)
 			proposalCtx = random.NewContext(proposalCtx, commit.Prop.Seed)
 			result, err := s.raftRepo.Apply(proposalCtx, commit.Prop)
 			commit.Done <- consensus.ProposalResult{Result: result, Error: err}

pkg/rid/store/memstore/identification_service_area.go

@@ -61,7 +61,7 @@ func (r *repo) InsertISA(ctx context.Context, isa *ridmodels.IdentificationServi
 		return nil, stacktrace.NewError("ISA with id %s already exists", isa.ID)
 	}
 
-	now := timestamp.MustGetRequestTimestamp(ctx)
+	now := timestamp.MustFromContext(ctx)
 
 	rec := isaRecordFromModel(isa, now)
 	r.state.ISAs[isa.ID] = rec
@@ -77,7 +77,7 @@ func (r *repo) UpdateISA(ctx context.Context, isa *ridmodels.IdentificationServi
 		return nil, nil
 	}
 
-	now := timestamp.MustGetRequestTimestamp(ctx)
+	now := timestamp.MustFromContext(ctx)
 
 	rec := isaRecordFromModel(isa, now)
 	rec.Owner = prev.Owner // It's not possible to update the owner of an ISA, this ensure it's to changed to a new value.

pkg/rid/store/memstore/identification_service_area_test.go

@@ -34,7 +34,7 @@ var (
 
 func TestStoreSearchISAs(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	cells := s2.CellUnion{
 		s2.CellID(17106221850767130624),
 		s2.CellID(17106221885126868992),
@@ -137,7 +137,7 @@ func TestStoreSearchISAs(t *testing.T) {
 
 func TestBadVersion(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	saOut1, err := repo.InsertISA(ctx, serviceArea)
@@ -159,7 +159,7 @@ func TestBadVersion(t *testing.T) {
 
 func TestStoreExpiredISA(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	saOut, err := repo.InsertISA(ctx, serviceArea)
@@ -194,7 +194,7 @@ func TestStoreExpiredISA(t *testing.T) {
 
 func TestStoreDeleteISAs(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	// Insert the ISA.
@@ -215,7 +215,7 @@ func TestStoreDeleteISAs(t *testing.T) {
 
 func TestStoreISAWithNoGeoData(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	endTime := fakeClock.Now().Add(24 * time.Hour)
@@ -230,7 +230,7 @@ func TestStoreISAWithNoGeoData(t *testing.T) {
 
 func TestListExpiredISAs(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	// Insert ISA with endtime 1 day from now
@@ -261,7 +261,7 @@ func TestListExpiredISAs(t *testing.T) {
 
 func TestListExpiredISAsWithEmptyWriter(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	// Insert ISA with endtime 1 day from now
@@ -294,7 +294,7 @@ func TestListExpiredISAsWithEmptyWriter(t *testing.T) {
 
 func TestStoreCountISAs(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	// Insert the ISA.

pkg/rid/store/memstore/snapshot_test.go

@@ -15,7 +15,7 @@ import (
 
 func TestSnapshotRoundTrip(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	src := setUpStore(t)
 	_, err := src.InsertISA(ctx, serviceArea)
 	require.NoError(t, err)
@@ -49,7 +49,7 @@ func TestSnapshotRoundTrip(t *testing.T) {
 
 func TestRestoreFromSnapshotReplacesState(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	src := setUpStore(t)
 	_, err := src.InsertISA(ctx, serviceArea)
 	require.NoError(t, err)

pkg/rid/store/memstore/store_test.go

@@ -30,7 +30,7 @@ func setUpStore(t *testing.T) *repo {
 
 func TestDatabaseEnsuresBeginsBeforeExpires(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	var (
@@ -50,7 +50,7 @@ func TestDatabaseEnsuresBeginsBeforeExpires(t *testing.T) {
 
 func TestCheckpointRestoreISA(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	_, err := repo.InsertISA(ctx, serviceArea)
@@ -76,7 +76,7 @@ func TestCheckpointRestoreISA(t *testing.T) {
 
 func TestCheckpointIsolatesNotificationIndex(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	sub, err := repo.InsertSubscription(ctx, subscriptionsPool[0].input)

pkg/rid/store/memstore/subscriptions.go

@@ -67,7 +67,7 @@ func (r *repo) InsertSubscription(ctx context.Context, s *ridmodels.Subscription
 		return nil, stacktrace.NewError("Subscription with id %s already exists", s.ID)
 	}
 
-	now := timestamp.MustGetRequestTimestamp(ctx)
+	now := timestamp.MustFromContext(ctx)
 
 	rec := subRecordFromModel(s, now)
 	r.state.Subscriptions[s.ID] = rec
@@ -83,7 +83,7 @@ func (r *repo) UpdateSubscription(ctx context.Context, s *ridmodels.Subscription
 		return nil, nil
 	}
 
-	now := timestamp.MustGetRequestTimestamp(ctx)
+	now := timestamp.MustFromContext(ctx)
 
 	rec := subRecordFromModel(s, now)
 	rec.Owner = prev.Owner // It's not possible to update the owner of a subscription, this ensure it's to changed to a new value.
@@ -137,7 +137,7 @@ func (r *repo) searchSubscriptions(ctx context.Context, cells s2.CellUnion, owne
 		return nil, stacktrace.NewErrorWithCode(dsserr.BadRequest, "no location provided")
 	}
 
-	now := timestamp.MustGetRequestTimestamp(ctx)
+	now := timestamp.MustFromContext(ctx)
 
 	var out []*ridmodels.Subscription
 	for rec := range r.liveSubscriptionsInCells(now, cells, owner) {
@@ -154,7 +154,7 @@ func (r *repo) searchSubscriptions(ctx context.Context, cells s2.CellUnion, owne
 // subscription in the given cells.
 func (r *repo) UpdateNotificationIdxsInCells(ctx context.Context, cells s2.CellUnion) ([]*ridmodels.Subscription, error) {
 
-	now := timestamp.MustGetRequestTimestamp(ctx)
+	now := timestamp.MustFromContext(ctx)
 
 	var out []*ridmodels.Subscription
 	for rec := range r.liveSubscriptionsInCells(now, cells, nil) {
@@ -166,7 +166,7 @@ func (r *repo) UpdateNotificationIdxsInCells(ctx context.Context, cells s2.CellU
 
 func (r *repo) MaxSubscriptionCountInCellsByOwner(ctx context.Context, cells s2.CellUnion, owner dssmodels.Owner) (int, error) {
 
-	now := timestamp.MustGetRequestTimestamp(ctx)
+	now := timestamp.MustFromContext(ctx)
 
 	want := cellSet(cells)
 	counts := make(map[s2.CellID]int, len(cells))