fix: avoid deleting peers on graceful close (#14165)

* fix: avoid deleting peers on graceful close

- Fixes an issue where a coordinator deletes all
  its peers on shutdown. This can cause disconnects
  whenever a coderd is redeployed.
This commit is contained in:
Jon Ayers
2024-08-14 15:16:08 -04:00
committed by GitHub
parent 6f1951e1c8
commit 4fc047954e
13 changed files with 330 additions and 104 deletions
+7
View File
@@ -3324,6 +3324,13 @@ func (q *querier) UpdateReplica(ctx context.Context, arg database.UpdateReplicaP
return q.db.UpdateReplica(ctx, arg)
}
func (q *querier) UpdateTailnetPeerStatusByCoordinator(ctx context.Context, arg database.UpdateTailnetPeerStatusByCoordinatorParams) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceTailnetCoordinator); err != nil {
return err
}
return q.db.UpdateTailnetPeerStatusByCoordinator(ctx, arg)
}
func (q *querier) UpdateTemplateACLByID(ctx context.Context, arg database.UpdateTemplateACLByIDParams) error {
fetch := func(ctx context.Context, arg database.UpdateTemplateACLByIDParams) (database.Template, error) {
return q.db.GetTemplateByID(ctx, arg.ID)
+5
View File
@@ -2152,6 +2152,11 @@ func (s *MethodTestSuite) TestTailnetFunctions() {
Asserts(rbac.ResourceTailnetCoordinator, policy.ActionCreate).
Errors(dbmem.ErrUnimplemented)
}))
s.Run("UpdateTailnetPeerStatusByCoordinator", s.Subtest(func(_ database.Store, check *expects) {
check.Args(database.UpdateTailnetPeerStatusByCoordinatorParams{}).
Asserts(rbac.ResourceTailnetCoordinator, policy.ActionUpdate).
Errors(dbmem.ErrUnimplemented)
}))
}
func (s *MethodTestSuite) TestDBCrypt() {
+4
View File
@@ -7917,6 +7917,10 @@ func (q *FakeQuerier) UpdateReplica(_ context.Context, arg database.UpdateReplic
return database.Replica{}, sql.ErrNoRows
}
func (*FakeQuerier) UpdateTailnetPeerStatusByCoordinator(context.Context, database.UpdateTailnetPeerStatusByCoordinatorParams) error {
return ErrUnimplemented
}
func (q *FakeQuerier) UpdateTemplateACLByID(_ context.Context, arg database.UpdateTemplateACLByIDParams) error {
if err := validateDatabaseType(arg); err != nil {
return err
+7
View File
@@ -2069,6 +2069,13 @@ func (m metricsStore) UpdateReplica(ctx context.Context, arg database.UpdateRepl
return replica, err
}
func (m metricsStore) UpdateTailnetPeerStatusByCoordinator(ctx context.Context, arg database.UpdateTailnetPeerStatusByCoordinatorParams) error {
start := time.Now()
r0 := m.s.UpdateTailnetPeerStatusByCoordinator(ctx, arg)
m.queryLatencies.WithLabelValues("UpdateTailnetPeerStatusByCoordinator").Observe(time.Since(start).Seconds())
return r0
}
func (m metricsStore) UpdateTemplateACLByID(ctx context.Context, arg database.UpdateTemplateACLByIDParams) error {
start := time.Now()
err := m.s.UpdateTemplateACLByID(ctx, arg)
+14
View File
@@ -4366,6 +4366,20 @@ func (mr *MockStoreMockRecorder) UpdateReplica(arg0, arg1 any) *gomock.Call {
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateReplica", reflect.TypeOf((*MockStore)(nil).UpdateReplica), arg0, arg1)
}
// UpdateTailnetPeerStatusByCoordinator mocks base method.
func (m *MockStore) UpdateTailnetPeerStatusByCoordinator(arg0 context.Context, arg1 database.UpdateTailnetPeerStatusByCoordinatorParams) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "UpdateTailnetPeerStatusByCoordinator", arg0, arg1)
ret0, _ := ret[0].(error)
return ret0
}
// UpdateTailnetPeerStatusByCoordinator indicates an expected call of UpdateTailnetPeerStatusByCoordinator.
func (mr *MockStoreMockRecorder) UpdateTailnetPeerStatusByCoordinator(arg0, arg1 any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTailnetPeerStatusByCoordinator", reflect.TypeOf((*MockStore)(nil).UpdateTailnetPeerStatusByCoordinator), arg0, arg1)
}
// UpdateTemplateACLByID mocks base method.
func (m *MockStore) UpdateTemplateACLByID(arg0 context.Context, arg1 database.UpdateTemplateACLByIDParams) error {
m.ctrl.T.Helper()
+10
View File
@@ -35,6 +35,7 @@ type options struct {
dumpOnFailure bool
returnSQLDB func(*sql.DB)
logger slog.Logger
url string
}
type Option func(*options)
@@ -59,6 +60,12 @@ func WithLogger(logger slog.Logger) Option {
}
}
func WithURL(u string) Option {
return func(o *options) {
o.url = u
}
}
func withReturnSQLDB(f func(*sql.DB)) Option {
return func(o *options) {
o.returnSQLDB = f
@@ -92,6 +99,9 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) {
ps := pubsub.NewInMemory()
if WillUsePostgres() {
connectionURL := os.Getenv("CODER_PG_CONNECTION_URL")
if connectionURL == "" && o.url != "" {
connectionURL = o.url
}
if connectionURL == "" {
var (
err error
+1
View File
@@ -418,6 +418,7 @@ type sqlcQuerier interface {
UpdateProvisionerJobWithCancelByID(ctx context.Context, arg UpdateProvisionerJobWithCancelByIDParams) error
UpdateProvisionerJobWithCompleteByID(ctx context.Context, arg UpdateProvisionerJobWithCompleteByIDParams) error
UpdateReplica(ctx context.Context, arg UpdateReplicaParams) (Replica, error)
UpdateTailnetPeerStatusByCoordinator(ctx context.Context, arg UpdateTailnetPeerStatusByCoordinatorParams) error
UpdateTemplateACLByID(ctx context.Context, arg UpdateTemplateACLByIDParams) error
UpdateTemplateAccessControlByID(ctx context.Context, arg UpdateTemplateAccessControlByIDParams) error
UpdateTemplateActiveVersionByID(ctx context.Context, arg UpdateTemplateActiveVersionByIDParams) error
+19
View File
@@ -7419,6 +7419,25 @@ func (q *sqlQuerier) GetTailnetTunnelPeerIDs(ctx context.Context, srcID uuid.UUI
return items, nil
}
const updateTailnetPeerStatusByCoordinator = `-- name: UpdateTailnetPeerStatusByCoordinator :exec
UPDATE
tailnet_peers
SET
status = $2
WHERE
coordinator_id = $1
`
type UpdateTailnetPeerStatusByCoordinatorParams struct {
CoordinatorID uuid.UUID `db:"coordinator_id" json:"coordinator_id"`
Status TailnetStatus `db:"status" json:"status"`
}
func (q *sqlQuerier) UpdateTailnetPeerStatusByCoordinator(ctx context.Context, arg UpdateTailnetPeerStatusByCoordinatorParams) error {
_, err := q.db.ExecContext(ctx, updateTailnetPeerStatusByCoordinator, arg.CoordinatorID, arg.Status)
return err
}
const upsertTailnetAgent = `-- name: UpsertTailnetAgent :one
INSERT INTO
tailnet_agents (
+8
View File
@@ -149,6 +149,14 @@ DO UPDATE SET
updated_at = now() at time zone 'utc'
RETURNING *;
-- name: UpdateTailnetPeerStatusByCoordinator :exec
UPDATE
tailnet_peers
SET
status = $2
WHERE
coordinator_id = $1;
-- name: DeleteTailnetPeer :one
DELETE
FROM tailnet_peers
+31 -27
View File
@@ -144,10 +144,6 @@ func newPGCoordInternal(
// signals when first heartbeat has been sent, so it's safe to start binding.
fHB := make(chan struct{})
// we need to arrange for the querier to stop _after_ the tunneler and binder, since we delete
// the coordinator when the querier stops (via the heartbeats). If the tunneler and binder are
// still running, they could run afoul of foreign key constraints.
querierCtx, querierCancel := context.WithCancel(dbauthz.As(context.Background(), pgCoordSubject))
c := &pgCoord{
ctx: ctx,
cancel: cancel,
@@ -163,18 +159,9 @@ func newPGCoordInternal(
handshaker: newHandshaker(ctx, logger, id, ps, rfhCh, fHB),
handshakerCh: rfhCh,
id: id,
querier: newQuerier(querierCtx, logger, id, ps, store, id, cCh, ccCh, numQuerierWorkers, fHB, clk),
querier: newQuerier(ctx, logger, id, ps, store, id, cCh, ccCh, numQuerierWorkers, fHB, clk),
closed: make(chan struct{}),
}
go func() {
// when the main context is canceled, or the coordinator closed, the binder, tunneler, and
// handshaker always eventually stop. Once they stop it's safe to cancel the querier context, which
// has the effect of deleting the coordinator from the database and ceasing heartbeats.
c.binder.workerWG.Wait()
c.tunneler.workerWG.Wait()
c.handshaker.workerWG.Wait()
querierCancel()
}()
logger.Info(ctx, "starting coordinator")
return c, nil
}
@@ -239,6 +226,9 @@ func (c *pgCoord) Close() error {
c.cancel()
c.closeOnce.Do(func() { close(c.closed) })
c.querier.wait()
c.binder.wait()
c.tunneler.workerWG.Wait()
c.handshaker.workerWG.Wait()
return nil
}
@@ -485,6 +475,7 @@ type binder struct {
workQ *workQ[bKey]
workerWG sync.WaitGroup
close chan struct{}
}
func newBinder(ctx context.Context,
@@ -502,6 +493,7 @@ func newBinder(ctx context.Context,
bindings: bindings,
latest: make(map[bKey]binding),
workQ: newWorkQ[bKey](ctx),
close: make(chan struct{}),
}
go b.handleBindings()
// add to the waitgroup immediately to avoid any races waiting for it before
@@ -513,6 +505,26 @@ func newBinder(ctx context.Context,
go b.worker()
}
}()
go func() {
defer close(b.close)
<-b.ctx.Done()
b.logger.Debug(b.ctx, "binder exiting, waiting for workers")
b.workerWG.Wait()
b.logger.Debug(b.ctx, "updating peers to lost")
ctx, cancel := context.WithTimeout(context.Background(), time.Second*15)
defer cancel()
err := b.store.UpdateTailnetPeerStatusByCoordinator(ctx, database.UpdateTailnetPeerStatusByCoordinatorParams{
CoordinatorID: b.coordinatorID,
Status: database.TailnetStatusLost,
})
if err != nil {
b.logger.Error(b.ctx, "update peer status to lost", slog.Error(err))
}
}()
return b
}
@@ -520,7 +532,7 @@ func (b *binder) handleBindings() {
for {
select {
case <-b.ctx.Done():
b.logger.Debug(b.ctx, "binder exiting", slog.Error(b.ctx.Err()))
b.logger.Debug(b.ctx, "binder exiting")
return
case bnd := <-b.bindings:
b.storeBinding(bnd)
@@ -632,6 +644,10 @@ func (b *binder) retrieveBinding(bk bKey) binding {
return bnd
}
func (b *binder) wait() {
<-b.close
}
// mapper tracks data sent to a peer, and sends updates based on changes read from the database.
type mapper struct {
ctx context.Context
@@ -1646,7 +1662,6 @@ func (h *heartbeats) sendBeats() {
// send an initial heartbeat so that other coordinators can start using our bindings right away.
h.sendBeat()
close(h.firstHeartbeat) // signal binder it can start writing
defer h.sendDelete()
tkr := h.clock.TickerFunc(h.ctx, HeartbeatPeriod, func() error {
h.sendBeat()
return nil
@@ -1677,17 +1692,6 @@ func (h *heartbeats) sendBeat() {
h.failedHeartbeats = 0
}
func (h *heartbeats) sendDelete() {
// here we don't want to use the main context, since it will have been canceled
ctx := dbauthz.As(context.Background(), pgCoordSubject)
err := h.store.DeleteCoordinator(ctx, h.self)
if err != nil {
h.logger.Error(h.ctx, "failed to send coordinator delete", slog.Error(err))
return
}
h.logger.Debug(h.ctx, "deleted coordinator")
}
func (h *heartbeats) cleanupLoop() {
defer h.wg.Done()
h.cleanup()
+1 -4
View File
@@ -396,10 +396,6 @@ func TestPGCoordinatorUnhealthy(t *testing.T) {
UpsertTailnetCoordinator(gomock.Any(), gomock.Any()).
Times(3).
Return(database.TailnetCoordinator{}, xerrors.New("badness"))
mStore.EXPECT().
DeleteCoordinator(gomock.Any(), gomock.Any()).
Times(1).
Return(nil)
// But, in particular we DO NOT want the coordinator to call DeleteTailnetPeer, as this is
// unnecessary and can spam the database. c.f. https://github.com/coder/coder/issues/12923
@@ -407,6 +403,7 @@ func TestPGCoordinatorUnhealthy(t *testing.T) {
mStore.EXPECT().CleanTailnetCoordinators(gomock.Any()).AnyTimes().Return(nil)
mStore.EXPECT().CleanTailnetLostPeers(gomock.Any()).AnyTimes().Return(nil)
mStore.EXPECT().CleanTailnetTunnels(gomock.Any()).AnyTimes().Return(nil)
mStore.EXPECT().UpdateTailnetPeerStatusByCoordinator(gomock.Any(), gomock.Any())
coordinator, err := newPGCoordInternal(ctx, logger, ps, mStore, mClock)
require.NoError(t, err)
+187 -64
View File
@@ -480,7 +480,7 @@ func TestPGCoordinatorSingle_SendsHeartbeats(t *testing.T) {
mu := sync.Mutex{}
heartbeats := []time.Time{}
unsub, err := ps.SubscribeWithErr(tailnet.EventHeartbeats, func(_ context.Context, msg []byte, err error) {
unsub, err := ps.SubscribeWithErr(tailnet.EventHeartbeats, func(_ context.Context, _ []byte, err error) {
assert.NoError(t, err)
mu.Lock()
defer mu.Unlock()
@@ -591,10 +591,10 @@ func TestPGCoordinatorDual_Mainline(t *testing.T) {
require.ErrorIs(t, err, io.EOF)
err = client21.recvErr(ctx, t)
require.ErrorIs(t, err, io.EOF)
assertEventuallyLost(ctx, t, store, agent2.id)
assertEventuallyLost(ctx, t, store, client21.id)
assertEventuallyLost(ctx, t, store, client22.id)
assertEventuallyNoAgents(ctx, t, store, agent2.id)
t.Logf("close coord1")
err = coord1.Close()
require.NoError(t, err)
// this closes agent1, client12, client11
@@ -604,6 +604,9 @@ func TestPGCoordinatorDual_Mainline(t *testing.T) {
require.ErrorIs(t, err, io.EOF)
err = client11.recvErr(ctx, t)
require.ErrorIs(t, err, io.EOF)
assertEventuallyLost(ctx, t, store, agent1.id)
assertEventuallyLost(ctx, t, store, client11.id)
assertEventuallyLost(ctx, t, store, client12.id)
// wait for all connections to close
err = agent1.close()
@@ -629,10 +632,6 @@ func TestPGCoordinatorDual_Mainline(t *testing.T) {
err = client22.close()
require.NoError(t, err)
client22.waitForClose(ctx, t)
assertEventuallyNoAgents(ctx, t, store, agent1.id)
assertEventuallyNoClientsForAgent(ctx, t, store, agent1.id)
assertEventuallyNoClientsForAgent(ctx, t, store, agent2.id)
}
// TestPGCoordinator_MultiCoordinatorAgent tests when a single agent connects to multiple coordinators.
@@ -746,7 +745,7 @@ func TestPGCoordinator_Unhealthy(t *testing.T) {
mStore.EXPECT().DeleteTailnetPeer(gomock.Any(), gomock.Any()).
AnyTimes().Return(database.DeleteTailnetPeerRow{}, nil)
mStore.EXPECT().DeleteAllTailnetTunnels(gomock.Any(), gomock.Any()).AnyTimes().Return(nil)
mStore.EXPECT().DeleteCoordinator(gomock.Any(), gomock.Any()).AnyTimes().Return(nil)
mStore.EXPECT().UpdateTailnetPeerStatusByCoordinator(gomock.Any(), gomock.Any())
uut, err := tailnet.NewPGCoord(ctx, logger, ps, mStore)
require.NoError(t, err)
@@ -811,7 +810,7 @@ func TestPGCoordinator_Node_Empty(t *testing.T) {
mStore.EXPECT().CleanTailnetCoordinators(gomock.Any()).AnyTimes().Return(nil)
mStore.EXPECT().CleanTailnetLostPeers(gomock.Any()).AnyTimes().Return(nil)
mStore.EXPECT().CleanTailnetTunnels(gomock.Any()).AnyTimes().Return(nil)
mStore.EXPECT().DeleteCoordinator(gomock.Any(), gomock.Any()).AnyTimes().Return(nil)
mStore.EXPECT().UpdateTailnetPeerStatusByCoordinator(gomock.Any(), gomock.Any()).Times(1)
uut, err := tailnet.NewPGCoord(ctx, logger, ps, mStore)
require.NoError(t, err)
@@ -871,51 +870,184 @@ func TestPGCoordinator_Lost(t *testing.T) {
agpltest.LostTest(ctx, t, coordinator)
}
func TestPGCoordinator_DeleteOnClose(t *testing.T) {
func TestPGCoordinator_NoDeleteOnClose(t *testing.T) {
t.Parallel()
if !dbtestutil.WillUsePostgres() {
t.Skip("test only with postgres")
}
store, ps := dbtestutil.NewDB(t)
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitSuperLong)
defer cancel()
ctrl := gomock.NewController(t)
mStore := dbmock.NewMockStore(ctrl)
ps := pubsub.NewInMemory()
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
coordinator, err := tailnet.NewPGCoord(ctx, logger, ps, store)
require.NoError(t, err)
defer coordinator.Close()
agent := newTestAgent(t, coordinator, "original")
defer agent.close()
agent.sendNode(&agpl.Node{PreferredDERP: 10})
client := newTestClient(t, coordinator, agent.id)
defer client.close()
// Simulate some traffic to generate
// a peer.
agentNodes := client.recvNodes(ctx, t)
require.Len(t, agentNodes, 1)
assert.Equal(t, 10, agentNodes[0].PreferredDERP)
client.sendNode(&agpl.Node{PreferredDERP: 11})
clientNodes := agent.recvNodes(ctx, t)
require.Len(t, clientNodes, 1)
assert.Equal(t, 11, clientNodes[0].PreferredDERP)
anode := coordinator.Node(agent.id)
require.NotNil(t, anode)
cnode := coordinator.Node(client.id)
require.NotNil(t, cnode)
err = coordinator.Close()
require.NoError(t, err)
assertEventuallyLost(ctx, t, store, agent.id)
assertEventuallyLost(ctx, t, store, client.id)
coordinator2, err := tailnet.NewPGCoord(ctx, logger, ps, store)
require.NoError(t, err)
defer coordinator2.Close()
anode = coordinator2.Node(agent.id)
require.NotNil(t, anode)
assert.Equal(t, 10, anode.PreferredDERP)
cnode = coordinator2.Node(client.id)
require.NotNil(t, cnode)
assert.Equal(t, 11, cnode.PreferredDERP)
}
// TestPGCoordinatorDual_FailedHeartbeat tests that peers
// disconnect from a coordinator when they are unhealthy,
// are marked as LOST (not DISCONNECTED), and can reconnect to
// a new coordinator and reestablish their tunnels.
func TestPGCoordinatorDual_FailedHeartbeat(t *testing.T) {
t.Parallel()
if !dbtestutil.WillUsePostgres() {
t.Skip("test only with postgres")
}
dburl, closeFn, err := dbtestutil.Open()
require.NoError(t, err)
t.Cleanup(closeFn)
store1, ps1, sdb1 := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithURL(dburl))
defer sdb1.Close()
store2, ps2, sdb2 := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithURL(dburl))
defer sdb2.Close()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitSuperLong)
t.Cleanup(cancel)
// We do this to avoid failing due errors related to the
// database connection being close.
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
upsertDone := make(chan struct{})
deleteCalled := make(chan struct{})
finishDelete := make(chan struct{})
mStore.EXPECT().UpsertTailnetCoordinator(gomock.Any(), gomock.Any()).
MinTimes(1).
Do(func(_ context.Context, _ uuid.UUID) { close(upsertDone) }).
Return(database.TailnetCoordinator{}, nil)
mStore.EXPECT().DeleteCoordinator(gomock.Any(), gomock.Any()).
Times(1).
Do(func(_ context.Context, _ uuid.UUID) {
close(deleteCalled)
<-finishDelete
}).
Return(nil)
// extra calls we don't particularly care about for this test
mStore.EXPECT().CleanTailnetCoordinators(gomock.Any()).AnyTimes().Return(nil)
mStore.EXPECT().CleanTailnetLostPeers(gomock.Any()).AnyTimes().Return(nil)
mStore.EXPECT().CleanTailnetTunnels(gomock.Any()).AnyTimes().Return(nil)
uut, err := tailnet.NewPGCoord(ctx, logger, ps, mStore)
// Create two coordinators, 1 for each peer.
c1, err := tailnet.NewPGCoord(ctx, logger, ps1, store1)
require.NoError(t, err)
testutil.RequireRecvCtx(ctx, t, upsertDone)
closeErr := make(chan error, 1)
go func() {
closeErr <- uut.Close()
}()
select {
case <-closeErr:
t.Fatal("close returned before DeleteCoordinator called")
case <-deleteCalled:
close(finishDelete)
err := testutil.RequireRecvCtx(ctx, t, closeErr)
require.NoError(t, err)
c2, err := tailnet.NewPGCoord(ctx, logger, ps2, store2)
require.NoError(t, err)
p1 := agpltest.NewPeer(ctx, t, c1, "peer1")
p2 := agpltest.NewPeer(ctx, t, c2, "peer2")
// Create a binding between the two.
p1.AddTunnel(p2.ID)
// Ensure that messages pass through.
p1.UpdateDERP(1)
p2.UpdateDERP(2)
p1.AssertEventuallyHasDERP(p2.ID, 2)
p2.AssertEventuallyHasDERP(p1.ID, 1)
// Close the underlying database connection to induce
// a heartbeat failure scenario and assert that
// we eventually disconnect from the coordinator.
err = sdb1.Close()
require.NoError(t, err)
p1.AssertEventuallyResponsesClosed()
p2.AssertEventuallyLost(p1.ID)
// This basically checks that peer2 had no update
// performed on their status since we are connected
// to coordinator2.
assertEventuallyStatus(ctx, t, store2, p2.ID, database.TailnetStatusOk)
// Connect peer1 to coordinator2.
p1.ConnectToCoordinator(ctx, c2)
// Reestablish binding.
p1.AddTunnel(p2.ID)
// Ensure messages still flow back and forth.
p1.AssertEventuallyHasDERP(p2.ID, 2)
p1.UpdateDERP(3)
p2.UpdateDERP(4)
p2.AssertEventuallyHasDERP(p1.ID, 3)
p1.AssertEventuallyHasDERP(p2.ID, 4)
// Make sure peer2 never got an update about peer1 disconnecting.
p2.AssertNeverUpdateKind(p1.ID, proto.CoordinateResponse_PeerUpdate_DISCONNECTED)
}
func TestPGCoordinatorDual_PeerReconnect(t *testing.T) {
t.Parallel()
if !dbtestutil.WillUsePostgres() {
t.Skip("test only with postgres")
}
store, ps := dbtestutil.NewDB(t)
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitSuperLong)
defer cancel()
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
// Create two coordinators, 1 for each peer.
c1, err := tailnet.NewPGCoord(ctx, logger, ps, store)
require.NoError(t, err)
c2, err := tailnet.NewPGCoord(ctx, logger, ps, store)
require.NoError(t, err)
p1 := agpltest.NewPeer(ctx, t, c1, "peer1")
p2 := agpltest.NewPeer(ctx, t, c2, "peer2")
// Create a binding between the two.
p1.AddTunnel(p2.ID)
// Ensure that messages pass through.
p1.UpdateDERP(1)
p2.UpdateDERP(2)
p1.AssertEventuallyHasDERP(p2.ID, 2)
p2.AssertEventuallyHasDERP(p1.ID, 1)
// Close coordinator1. Now we will check that we
// never send a DISCONNECTED update.
err = c1.Close()
require.NoError(t, err)
p1.AssertEventuallyResponsesClosed()
p2.AssertEventuallyLost(p1.ID)
// This basically checks that peer2 had no update
// performed on their status since we are connected
// to coordinator2.
assertEventuallyStatus(ctx, t, store, p2.ID, database.TailnetStatusOk)
// Connect peer1 to coordinator2.
p1.ConnectToCoordinator(ctx, c2)
// Reestablish binding.
p1.AddTunnel(p2.ID)
// Ensure messages still flow back and forth.
p1.AssertEventuallyHasDERP(p2.ID, 2)
p1.UpdateDERP(3)
p2.UpdateDERP(4)
p2.AssertEventuallyHasDERP(p1.ID, 3)
p1.AssertEventuallyHasDERP(p2.ID, 4)
// Make sure peer2 never got an update about peer1 disconnecting.
p2.AssertNeverUpdateKind(p1.ID, proto.CoordinateResponse_PeerUpdate_DISCONNECTED)
}
type testConn struct {
@@ -1056,21 +1188,7 @@ func assertNeverHasDERPs(ctx context.Context, t *testing.T, c *testConn, expecte
}
}
func assertEventuallyNoAgents(ctx context.Context, t *testing.T, store database.Store, agentID uuid.UUID) {
t.Helper()
assert.Eventually(t, func() bool {
agents, err := store.GetTailnetPeers(ctx, agentID)
if xerrors.Is(err, sql.ErrNoRows) {
return true
}
if err != nil {
t.Fatal(err)
}
return len(agents) == 0
}, testutil.WaitShort, testutil.IntervalFast)
}
func assertEventuallyLost(ctx context.Context, t *testing.T, store database.Store, agentID uuid.UUID) {
func assertEventuallyStatus(ctx context.Context, t *testing.T, store database.Store, agentID uuid.UUID, status database.TailnetStatus) {
t.Helper()
assert.Eventually(t, func() bool {
peers, err := store.GetTailnetPeers(ctx, agentID)
@@ -1081,7 +1199,7 @@ func assertEventuallyLost(ctx context.Context, t *testing.T, store database.Stor
t.Fatal(err)
}
for _, peer := range peers {
if peer.Status == database.TailnetStatusOk {
if peer.Status != status {
return false
}
}
@@ -1089,6 +1207,11 @@ func assertEventuallyLost(ctx context.Context, t *testing.T, store database.Stor
}, testutil.WaitShort, testutil.IntervalFast)
}
func assertEventuallyLost(ctx context.Context, t *testing.T, store database.Store, agentID uuid.UUID) {
t.Helper()
assertEventuallyStatus(ctx, t, store, agentID, database.TailnetStatusLost)
}
func assertEventuallyNoClientsForAgent(ctx context.Context, t *testing.T, store database.Store, agentID uuid.UUID) {
t.Helper()
assert.Eventually(t, func() bool {
+36 -9
View File
@@ -19,18 +19,24 @@ type PeerStatus struct {
}
type Peer struct {
ctx context.Context
cancel context.CancelFunc
t testing.TB
ID uuid.UUID
name string
resps <-chan *proto.CoordinateResponse
reqs chan<- *proto.CoordinateRequest
peers map[uuid.UUID]PeerStatus
ctx context.Context
cancel context.CancelFunc
t testing.TB
ID uuid.UUID
name string
resps <-chan *proto.CoordinateResponse
reqs chan<- *proto.CoordinateRequest
peers map[uuid.UUID]PeerStatus
peerUpdates map[uuid.UUID][]*proto.CoordinateResponse_PeerUpdate
}
func NewPeer(ctx context.Context, t testing.TB, coord tailnet.CoordinatorV2, name string, id ...uuid.UUID) *Peer {
p := &Peer{t: t, name: name, peers: make(map[uuid.UUID]PeerStatus)}
p := &Peer{
t: t,
name: name,
peers: make(map[uuid.UUID]PeerStatus),
peerUpdates: make(map[uuid.UUID][]*proto.CoordinateResponse_PeerUpdate),
}
p.ctx, p.cancel = context.WithCancel(ctx)
if len(id) > 1 {
t.Fatal("too many")
@@ -45,6 +51,12 @@ func NewPeer(ctx context.Context, t testing.TB, coord tailnet.CoordinatorV2, nam
return p
}
func (p *Peer) ConnectToCoordinator(ctx context.Context, c tailnet.CoordinatorV2) {
p.t.Helper()
p.reqs, p.resps = c.Coordinate(ctx, p.ID, p.name, tailnet.SingleTailnetCoordinateeAuth{})
}
func (p *Peer) AddTunnel(other uuid.UUID) {
p.t.Helper()
req := &proto.CoordinateRequest{AddTunnel: &proto.CoordinateRequest_Tunnel{Id: tailnet.UUIDToByteSlice(other)}}
@@ -180,6 +192,19 @@ func (p *Peer) AssertEventuallyGetsError(match string) {
}
}
// AssertNeverUpdateKind asserts that we have not received
// any updates on the provided peer for the provided kind.
func (p *Peer) AssertNeverUpdateKind(peer uuid.UUID, kind proto.CoordinateResponse_PeerUpdate_Kind) {
p.t.Helper()
updates, ok := p.peerUpdates[peer]
assert.True(p.t, ok, "expected updates for peer %s", peer)
for _, update := range updates {
assert.NotEqual(p.t, kind, update.Kind, update)
}
}
var responsesClosed = xerrors.New("responses closed")
func (p *Peer) handleOneResp() error {
@@ -198,6 +223,8 @@ func (p *Peer) handleOneResp() error {
if err != nil {
return err
}
p.peerUpdates[id] = append(p.peerUpdates[id], update)
switch update.Kind {
case proto.CoordinateResponse_PeerUpdate_NODE, proto.CoordinateResponse_PeerUpdate_LOST:
peer := p.peers[id]