Skip to content

Commit ca2e6e7

Browse files
the-gluthe-glu
committed
[raft/memstore] Add aux memstore
1 parent 2bb37c0 commit ca2e6e7

3 files changed

Lines changed: 216 additions & 6 deletions

File tree

‎pkg/aux_/store/memstore/dss.go‎

Lines changed: 65 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,24 +2,85 @@ package memstore
22

33
import (
44
"context"
5+
"database/sql"
6+
"time"
57

68
auxmodels "github.com/interuss/dss/pkg/aux_/models"
79
dsserr "github.com/interuss/dss/pkg/errors"
810
"github.com/interuss/stacktrace"
911
)
1012

1113
func (r *repo) SaveOwnMetadata(_ context.Context, locality string, publicEndpoint string) error {
12-
return stacktrace.NewErrorWithCode(dsserr.NotImplemented, "SaveOwnMetadata not implemented for memstore")
14+
if locality == "" {
15+
return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Locality not set")
16+
}
17+
if publicEndpoint == "" {
18+
return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Public endpoint not set")
19+
}
20+
21+
r.participants[locality] = &participant{
22+
publicEndpoint: publicEndpoint,
23+
updatedAt: time.Now(),
24+
}
25+
return nil
1326
}
1427

1528
func (r *repo) GetDSSMetadata(_ context.Context) ([]*auxmodels.DSSMetadata, error) {
16-
return nil, stacktrace.NewErrorWithCode(dsserr.NotImplemented, "GetDSSMetadata not implemented for memstore")
29+
metadata := make([]*auxmodels.DSSMetadata, 0, len(r.participants))
30+
for locality, p := range r.participants {
31+
updatedAt := p.updatedAt
32+
m := &auxmodels.DSSMetadata{
33+
Locality: locality,
34+
PublicEndpoint: p.publicEndpoint,
35+
UpdatedAt: &updatedAt,
36+
}
37+
38+
// Find the latest heartbeat across all sources for this locality.
39+
var latest auxmodels.Heartbeat
40+
found := false
41+
for key, hb := range r.heartbeats {
42+
if key.locality != locality {
43+
continue
44+
}
45+
if !found || hb.Timestamp.After(*latest.Timestamp) {
46+
latest = hb
47+
found = true
48+
}
49+
}
50+
51+
if found {
52+
m.LatestTimestamp.Source = sql.NullString{String: latest.Source, Valid: true}
53+
m.LatestTimestamp.Timestamp = latest.Timestamp
54+
m.LatestTimestamp.NextHeartbeatExpectedBefore = latest.NextHeartbeatExpectedBefore
55+
m.LatestTimestamp.Reporter = sql.NullString{String: latest.Reporter, Valid: true}
56+
}
57+
58+
metadata = append(metadata, m)
59+
}
60+
return metadata, nil
1761
}
1862

1963
func (r *repo) RecordHeartbeat(_ context.Context, heartbeat auxmodels.Heartbeat) error {
20-
return stacktrace.NewErrorWithCode(dsserr.NotImplemented, "RecordHeartbeat not implemented for memstore")
64+
if heartbeat.Locality == "" {
65+
return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Locality not set")
66+
}
67+
if heartbeat.Source == "" {
68+
return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Source not set")
69+
}
70+
71+
if heartbeat.Timestamp == nil {
72+
now := time.Now()
73+
heartbeat.Timestamp = &now
74+
}
75+
76+
if heartbeat.NextHeartbeatExpectedBefore != nil && heartbeat.NextHeartbeatExpectedBefore.Before(*heartbeat.Timestamp) {
77+
return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Cannot expect the timestamp of the next heartbeat before the timestamp of the new heartbeat")
78+
}
79+
80+
r.heartbeats[heartbeatKey{locality: heartbeat.Locality, source: heartbeat.Source}] = heartbeat
81+
return nil
2182
}
2283

2384
func (r *repo) GetDSSAirspaceRepresentationID(_ context.Context) (string, error) {
24-
return "", stacktrace.NewErrorWithCode(dsserr.NotImplemented, "GetDSSAirspaceRepresentationID not implemented for memstore")
85+
return "", stacktrace.NewErrorWithCode(dsserr.NotImplemented, "GetDSSAirspaceRepresentationID not implementable for memstore")
2586
}

‎pkg/aux_/store/memstore/dss_test.go‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
package memstore
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
auxmodels "github.com/interuss/dss/pkg/aux_/models"
9+
dsserr "github.com/interuss/dss/pkg/errors"
10+
"github.com/interuss/stacktrace"
11+
"github.com/stretchr/testify/require"
12+
)
13+
14+
func TestSaveOwnMetadataValidation(t *testing.T) {
15+
ctx := context.Background()
16+
r := newRepo()
17+
18+
require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(r.SaveOwnMetadata(ctx, "", "https://example.com")))
19+
require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(r.SaveOwnMetadata(ctx, "dss-1", "")))
20+
}
21+
22+
func TestSaveOwnMetadataRoundTrip(t *testing.T) {
23+
ctx := context.Background()
24+
r := newRepo()
25+
26+
require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
27+
28+
md, err := r.GetDSSMetadata(ctx)
29+
require.NoError(t, err)
30+
31+
require.Len(t, md, 1)
32+
require.Equal(t, "dss-1", md[0].Locality)
33+
require.Equal(t, "https://example.com", md[0].PublicEndpoint)
34+
require.NotNil(t, md[0].UpdatedAt)
35+
36+
// No heartbeat recorded yet.
37+
require.False(t, md[0].LatestTimestamp.Source.Valid)
38+
require.Nil(t, md[0].LatestTimestamp.Timestamp)
39+
}
40+
41+
func TestSaveOwnMetadataUpsert(t *testing.T) {
42+
ctx := context.Background()
43+
r := newRepo()
44+
45+
require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://old.example.com"))
46+
require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://new.example.com"))
47+
48+
md, err := r.GetDSSMetadata(ctx)
49+
require.NoError(t, err)
50+
51+
require.Len(t, md, 1)
52+
require.Equal(t, "https://new.example.com", md[0].PublicEndpoint)
53+
}
54+
55+
func TestRecordHeartbeatValidation(t *testing.T) {
56+
ctx := context.Background()
57+
r := newRepo()
58+
59+
require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Source: "source1"})))
60+
require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1"})))
61+
62+
ts := time.Now()
63+
before := ts.Add(-time.Minute)
64+
err := r.RecordHeartbeat(ctx, auxmodels.Heartbeat{
65+
Locality: "dss-1",
66+
Source: "source1",
67+
Timestamp: &ts,
68+
NextHeartbeatExpectedBefore: &before,
69+
})
70+
71+
require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(err))
72+
}
73+
74+
func TestRecordHeartbeatDefaultsTimestamp(t *testing.T) {
75+
ctx := context.Background()
76+
r := newRepo()
77+
78+
require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
79+
require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source1"}))
80+
81+
md, err := r.GetDSSMetadata(ctx)
82+
require.NoError(t, err)
83+
84+
require.Len(t, md, 1)
85+
require.True(t, md[0].LatestTimestamp.Source.Valid)
86+
require.NotNil(t, md[0].LatestTimestamp.Timestamp)
87+
}
88+
89+
func TestGetDSSMetadataPicksLatestHeartbeat(t *testing.T) {
90+
ctx := context.Background()
91+
r := newRepo()
92+
93+
require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
94+
95+
older := time.Now().Add(-time.Hour)
96+
newer := time.Now()
97+
require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source1", Timestamp: &older, Reporter: "uss1"}))
98+
require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source2", Timestamp: &newer, Reporter: "uss2"}))
99+
100+
md, err := r.GetDSSMetadata(ctx)
101+
require.NoError(t, err)
102+
103+
require.Len(t, md, 1)
104+
require.True(t, md[0].LatestTimestamp.Timestamp.Equal(newer))
105+
require.Equal(t, "source2", md[0].LatestTimestamp.Source.String)
106+
require.Equal(t, "uss2", md[0].LatestTimestamp.Reporter.String)
107+
}
108+
109+
func TestGetDSSMetadataUpdatesHeartbeatPerSource(t *testing.T) {
110+
ctx := context.Background()
111+
r := newRepo()
112+
113+
require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
114+
115+
first := time.Now().Add(-time.Hour)
116+
second := time.Now()
117+
require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source1", Timestamp: &first}))
118+
require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source1", Timestamp: &second}))
119+
120+
md, err := r.GetDSSMetadata(ctx)
121+
require.NoError(t, err)
122+
123+
require.Len(t, md, 1)
124+
require.True(t, md[0].LatestTimestamp.Timestamp.Equal(second))
125+
}

‎pkg/aux_/store/memstore/store.go‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,41 @@ package memstore
22

33
import (
44
"context"
5+
"time"
56

7+
auxmodels "github.com/interuss/dss/pkg/aux_/models"
68
"github.com/interuss/dss/pkg/aux_/repos"
79
"github.com/interuss/dss/pkg/memstore"
810
"go.uber.org/zap"
911
)
1012

1113
// repo is a full implementation of aux_.repos.Repository for memory-based storage.
12-
type repo struct{}
14+
type repo struct {
15+
// participants holds pool participants metadata, keyed by locality.
16+
participants map[string]*participant
17+
// heartbeats holds the latest heartbeat per (locality, source).
18+
heartbeats map[heartbeatKey]auxmodels.Heartbeat
19+
}
20+
21+
type participant struct {
22+
publicEndpoint string
23+
updatedAt time.Time
24+
}
25+
26+
type heartbeatKey struct {
27+
locality string
28+
source string
29+
}
30+
31+
func newRepo() *repo {
32+
return &repo{
33+
participants: map[string]*participant{},
34+
heartbeats: map[heartbeatKey]auxmodels.Heartbeat{},
35+
}
36+
}
1337

1438
func Init(ctx context.Context, logger *zap.Logger) (*memstore.Store[repos.Repository], error) {
15-
return memstore.Init(ctx, logger, "aux_", &repo{})
39+
return memstore.Init(ctx, logger, "aux_", newRepo())
1640
}
1741

1842
func (r *repo) GetRepo() repos.Repository { return r }

0 commit comments

Comments
 (0)