May 24, 2024 · May 23, 2024 · May 23, 2024
diff --git a/enterprise/tailnet/pgcoord.go b/enterprise/tailnet/pgcoord.go
 closed:           make(chan struct{}),
 }
 go func() {
 // when the main context is canceled, or the coordinator closed, the binder andtunneler
 // always eventually stop.  Once they stop it's safe to cancel the querier context, which
 // when the main context is canceled, or the coordinator closed, the binder,tunneler, and
 //handshakeralways 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")
 c.logger.Info(c.ctx, "closing coordinator")
 c.cancel()
 c.closeOnce.Do(func() { close(c.closed) })
 c.querier.wait()
 return nil
 }


 workQ *workQ[querierWorkKey]

 wg sync.WaitGroup

 heartbeats *heartbeats
 updates    <-chan hbUpdate

 }
 q.subscribe()

 q.wg.Add(2 + numWorkers)
 go func() {
 <-firstHeartbeat
 go q.handleIncoming()
 return q
 }

 func (q *querier) wait() {
 q.wg.Wait()
 q.heartbeats.wg.Wait()
 }

 func (q *querier) handleIncoming() {
 defer q.wg.Done()
 for {
 select {
 case <-q.ctx.Done():
 }

 func (q *querier) worker() {
 defer q.wg.Done()
 eb := backoff.NewExponentialBackOff()
 eb.MaxElapsedTime = 0 // retry indefinitely
 eb.MaxInterval = dbMaxBackoff
 }

 func (q *querier) handleUpdates() {
 defer q.wg.Done()
 for {
 select {
 case <-q.ctx.Done():
 coordinators map[uuid.UUID]time.Time
 timer        *time.Timer

 wg sync.WaitGroup

 // overwritten in tests, but otherwise constant
 cleanupPeriod time.Duration
 }
 coordinators:   make(map[uuid.UUID]time.Time),
 cleanupPeriod:  cleanupPeriod,
 }
 h.wg.Add(3)
 go h.subscribe()
 go h.sendBeats()
 go h.cleanupLoop()
 }

 func (h *heartbeats) subscribe() {
 defer h.wg.Done()
 eb := backoff.NewExponentialBackOff()
 eb.MaxElapsedTime = 0 // retry indefinitely
 eb.MaxInterval = dbMaxBackoff
 }

 func (h *heartbeats) sendBeats() {
 defer h.wg.Done()
 // 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
 }

 func (h *heartbeats) cleanupLoop() {
 defer h.wg.Done()
 h.cleanup()
 tkr := time.NewTicker(h.cleanupPeriod)
 defer tkr.Stop()
diff --git a/enterprise/tailnet/pgcoord_internal_test.go b/enterprise/tailnet/pgcoord_internal_test.go
 store:         mStore,
 cleanupPeriod: time.Millisecond,
 }
 uut.wg.Add(1)
 go uut.cleanupLoop()

 for i := 0; i < 6; i++ {
diff --git a/enterprise/tailnet/pgcoord_test.go b/enterprise/tailnet/pgcoord_test.go
 agpltest.LostTest(ctx, t, coordinator)
 }

 func TestPGCoordinator_DeleteOnClose(t *testing.T) {
 t.Parallel()

 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, &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)
 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)
 }
 }

 type testConn struct {
 ws, serverWS net.Conn
 nodeChan     chan []*agpl.Node
Original file line number	Diff line number	Diff line change
Expand Up		@@ -161,11 +161,12 @@ func newPGCoordInternal(
		closed: make(chan struct{}),
		}
		go func() {
		// when the main context is canceled, or the coordinator closed, the binder andtunneler
		// always eventually stop. Once they stop it's safe to cancel the querier context, which
		// when the main context is canceled, or the coordinator closed, the binder,tunneler, and
		//handshakeralways 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()
Copy link Contributor coadlerMay 23, 2024 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. Ah, thanks for adding this.
		querierCancel()
		}()
		logger.Info(ctx, "starting coordinator")
Expand DownExpand Up		@@ -231,6 +232,7 @@ func (c *pgCoord) Close() error {
		c.logger.Info(c.ctx, "closing coordinator")
		c.cancel()
		c.closeOnce.Do(func() { close(c.closed) })
		c.querier.wait()
		return nil
		}

Expand DownExpand Up		@@ -795,6 +797,8 @@ type querier struct {

		workQ *workQ[querierWorkKey]

		wg sync.WaitGroup

		heartbeats *heartbeats
		updates <-chan hbUpdate

Expand DownExpand Up		@@ -831,6 +835,7 @@ func newQuerier(ctx context.Context,
		}
		q.subscribe()

		q.wg.Add(2 + numWorkers)
		go func() {
		<-firstHeartbeat
		go q.handleIncoming()
Expand All		@@ -842,7 +847,13 @@ func newQuerier(ctx context.Context,
		return q
		}

		func (q *querier) wait() {
		q.wg.Wait()
		q.heartbeats.wg.Wait()
		}

		func (q *querier) handleIncoming() {
		defer q.wg.Done()
		for {
		select {
		case <-q.ctx.Done():
Expand DownExpand Up		@@ -919,6 +930,7 @@ func (q querier) cleanupConn(c connIO) {
		}

		func (q *querier) worker() {
		defer q.wg.Done()
		eb := backoff.NewExponentialBackOff()
		eb.MaxElapsedTime = 0 // retry indefinitely
		eb.MaxInterval = dbMaxBackoff
Expand DownExpand Up		@@ -1204,6 +1216,7 @@ func (q *querier) resyncPeerMappings() {
		}

		func (q *querier) handleUpdates() {
		defer q.wg.Done()
		for {
		select {
		case <-q.ctx.Done():
Expand DownExpand Up		@@ -1451,6 +1464,8 @@ type heartbeats struct {
		coordinators map[uuid.UUID]time.Time
		timer *time.Timer

		wg sync.WaitGroup

		// overwritten in tests, but otherwise constant
		cleanupPeriod time.Duration
		}
Expand All		@@ -1472,6 +1487,7 @@ func newHeartbeats(
		coordinators: make(map[uuid.UUID]time.Time),
		cleanupPeriod: cleanupPeriod,
		}
		h.wg.Add(3)
		go h.subscribe()
		go h.sendBeats()
		go h.cleanupLoop()
Expand DownExpand Up		@@ -1502,6 +1518,7 @@ func (h *heartbeats) filter(mappings []mapping) []mapping {
		}

		func (h *heartbeats) subscribe() {
		defer h.wg.Done()
		eb := backoff.NewExponentialBackOff()
		eb.MaxElapsedTime = 0 // retry indefinitely
		eb.MaxInterval = dbMaxBackoff
Expand DownExpand Up		@@ -1611,6 +1628,7 @@ func (h *heartbeats) checkExpiry() {
		}

		func (h *heartbeats) sendBeats() {
		defer h.wg.Done()
		// 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
Expand DownExpand Up		@@ -1662,6 +1680,7 @@ func (h *heartbeats) sendDelete() {
		}

		func (h *heartbeats) cleanupLoop() {
		defer h.wg.Done()
		h.cleanup()
		tkr := time.NewTicker(h.cleanupPeriod)
		defer tkr.Stop()
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -66,6 +66,7 @@ func TestHeartbeats_Cleanup(t *testing.T) {
		store: mStore,
		cleanupPeriod: time.Millisecond,
		}
		uut.wg.Add(1)
		go uut.cleanupLoop()

		for i := 0; i < 6; i++ {
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -864,6 +864,53 @@ func TestPGCoordinator_Lost(t *testing.T) {
		agpltest.LostTest(ctx, t, coordinator)
		}

		func TestPGCoordinator_DeleteOnClose(t *testing.T) {
		t.Parallel()

		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, &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)
		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)
		}
		}

		type testConn struct {
		ws, serverWS net.Conn
		nodeChan chan []*agpl.Node
Expand Down