Avoid race by exposing number of pending updates

dannykopping · dannykopping · commit 613e07444884 · 2024-07-04T12:13:51.000+02:00
Signed-off-by: Danny Kopping &lt;danny@coder.com&gt;
diff --git a/coderd/notifications/manager.go b/coderd/notifications/manager.go
@@ -44,6 +44,8 @@ type Manager struct {
 	notifier *notifier
 	handlers map[database.NotificationMethod]Handler
 
+	success, failure chan dispatchResult
+
 	runOnce  sync.Once
 	stopOnce sync.Once
 	stop     chan any
@@ -67,6 +69,15 @@ func NewManager(cfg codersdk.NotificationsConfig, store Store, log slog.Logger)
 		cfg:   cfg,
 		store: store,
 
+		// Buffer successful/failed notification dispatches in memory to reduce load on the store.
+		//
+		// We keep separate buffered for success/failure right now because the bulk updates are already a bit janky,
+		// see BulkMarkNotificationMessagesSent/BulkMarkNotificationMessagesFailed. If we had the ability to batch updates,
+		// like is offered in https://docs.sqlc.dev/en/stable/reference/query-annotations.html#batchmany, we'd have a cleaner
+		// approach to this - but for now this will work fine.
+		success: make(chan dispatchResult, cfg.StoreSyncBufferSize),
+		failure: make(chan dispatchResult, cfg.StoreSyncBufferSize),
+
 		stop: make(chan any),
 		done: make(chan any),
 
@@ -123,23 +134,12 @@ func (m *Manager) loop(ctx context.Context) error {
 	default:
 	}
 
-	var (
-		// Buffer successful/failed notification dispatches in memory to reduce load on the store.
-		//
-		// We keep separate buffered for success/failure right now because the bulk updates are already a bit janky,
-		// see BulkMarkNotificationMessagesSent/BulkMarkNotificationMessagesFailed. If we had the ability to batch updates,
-		// like is offered in https://docs.sqlc.dev/en/stable/reference/query-annotations.html#batchmany, we'd have a cleaner
-		// approach to this - but for now this will work fine.
-		success = make(chan dispatchResult, m.cfg.StoreSyncBufferSize)
-		failure = make(chan dispatchResult, m.cfg.StoreSyncBufferSize)
-	)
-
 	var eg errgroup.Group
 
 	// Create a notifier to run concurrently, which will handle dequeueing and dispatching notifications.
 	m.notifier = newNotifier(m.cfg, uuid.New(), m.log, m.store, m.handlers)
 	eg.Go(func() error {
-		return m.notifier.run(ctx, success, failure)
+		return m.notifier.run(ctx, m.success, m.failure)
 	})
 
 	// Periodically flush notification state changes to the store.
@@ -162,21 +162,21 @@ func (m *Manager) loop(ctx context.Context) error {
 				// TODO: mention the above tradeoff in documentation.
 				m.log.Warn(ctx, "exiting ungracefully", slog.Error(ctx.Err()))
 
-				if len(success)+len(failure) > 0 {
+				if len(m.success)+len(m.failure) > 0 {
 					m.log.Warn(ctx, "content canceled with pending updates in buffer, these messages will be sent again after lease expires",
-						slog.F("success_count", len(success)), slog.F("failure_count", len(failure)))
+						slog.F("success_count", len(m.success)), slog.F("failure_count", len(m.failure)))
 				}
 				return ctx.Err()
 			case <-m.stop:
-				if len(success)+len(failure) > 0 {
+				if len(m.success)+len(m.failure) > 0 {
 					m.log.Warn(ctx, "flushing buffered updates before stop",
-						slog.F("success_count", len(success)), slog.F("failure_count", len(failure)))
-					m.bulkUpdate(ctx, success, failure)
+						slog.F("success_count", len(m.success)), slog.F("failure_count", len(m.failure)))
+					m.bulkUpdate(ctx)
 					m.log.Warn(ctx, "flushing updates done")
 				}
 				return nil
 			case <-tick.C:
-				m.bulkUpdate(ctx, success, failure)
+				m.bulkUpdate(ctx)
 			}
 		}
 	})
@@ -188,16 +188,22 @@ func (m *Manager) loop(ctx context.Context) error {
 	return err
 }
 
+// BufferedUpdatesCount returns the number of buffered updates which are currently waiting to be flushed to the store.
+// The returned values are for success & failure, respectively.
+func (m *Manager) BufferedUpdatesCount() (success int, failure int) {
+	return len(m.success), len(m.failure)
+}
+
 // bulkUpdate updates messages in the store based on the given successful and failed message dispatch results.
-func (m *Manager) bulkUpdate(ctx context.Context, success, failure <-chan dispatchResult) {
+func (m *Manager) bulkUpdate(ctx context.Context) {
 	select {
 	case <-ctx.Done():
 		return
 	default:
 	}
 
-	nSuccess := len(success)
-	nFailure := len(failure)
+	nSuccess := len(m.success)
+	nFailure := len(m.failure)
 
 	// Nothing to do.
 	if nSuccess+nFailure == 0 {
@@ -217,12 +223,12 @@ func (m *Manager) bulkUpdate(ctx context.Context, success, failure <-chan dispat
 	// will be processed on the next bulk update.
 
 	for i := 0; i < nSuccess; i++ {
-		res := <-success
+		res := <-m.success
 		successParams.IDs = append(successParams.IDs, res.msg)
 		successParams.SentAts = append(successParams.SentAts, res.ts)
 	}
 	for i := 0; i < nFailure; i++ {
-		res := <-failure
+		res := <-m.failure
 
 		status := database.NotificationMessageStatusPermanentFailure
 		if res.retryable {
diff --git a/coderd/notifications/manager_test.go b/coderd/notifications/manager_test.go
@@ -5,7 +5,9 @@ import (
 	"encoding/json"
 	"sync/atomic"
 	"testing"
+	"time"
 
+	"github.com/coder/serpent"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -31,11 +33,14 @@ func TestBufferedUpdates(t *testing.T) {
 	if !dbtestutil.WillUsePostgres() {
 		t.Skip("This test requires postgres")
 	}
+
 	ctx, logger, db := setup(t)
 	interceptor := &bulkUpdateInterceptor{Store: db}
-
 	santa := &santaHandler{}
+
 	cfg := defaultNotificationsConfig(database.NotificationMethodSmtp)
+	cfg.StoreSyncInterval = serpent.Duration(time.Hour) // Ensure we don't sync the store automatically.
+
 	mgr, err := notifications.NewManager(cfg, interceptor, logger.Named("notifications-manager"))
 	require.NoError(t, err)
 	mgr.WithHandlers(map[database.NotificationMethod]notifications.Handler{
@@ -47,20 +52,34 @@ func TestBufferedUpdates(t *testing.T) {
 	user := dbgen.User(t, db, database.User{})
 
 	// given
-	_, err = enq.Enqueue(ctx, user.ID, notifications.TemplateWorkspaceDeleted, map[string]string{"nice": "true"}, "")
+	_, err = enq.Enqueue(ctx, user.ID, notifications.TemplateWorkspaceDeleted, map[string]string{"nice": "true"}, "") // Will succeed.
 	require.NoError(t, err)
-	_, err = enq.Enqueue(ctx, user.ID, notifications.TemplateWorkspaceDeleted, map[string]string{"nice": "true"}, "")
+	_, err = enq.Enqueue(ctx, user.ID, notifications.TemplateWorkspaceDeleted, map[string]string{"nice": "true"}, "") // Will succeed.
 	require.NoError(t, err)
-	_, err = enq.Enqueue(ctx, user.ID, notifications.TemplateWorkspaceDeleted, map[string]string{"nice": "false"}, "")
+	_, err = enq.Enqueue(ctx, user.ID, notifications.TemplateWorkspaceDeleted, map[string]string{"nice": "false"}, "") // Will fail.
 	require.NoError(t, err)
 
 	// when
 	mgr.Run(ctx)
 
 	// then
 
+	const (
+		expectedSuccess = 2
+		expectedFailure = 1
+	)
+
 	// Wait for messages to be dispatched.
-	require.Eventually(t, func() bool { return santa.naughty.Load() == 1 && santa.nice.Load() == 2 }, testutil.WaitMedium, testutil.IntervalFast)
+	require.Eventually(t, func() bool {
+		return santa.naughty.Load() == expectedFailure &&
+			santa.nice.Load() == expectedSuccess
+	}, testutil.WaitMedium, testutil.IntervalFast)
+
+	// Wait for the expected number of buffered updates to be accumulated.
+	require.Eventually(t, func() bool {
+		success, failure := mgr.BufferedUpdatesCount()
+		return success == expectedSuccess && failure == expectedFailure
+	}, testutil.WaitShort, testutil.IntervalFast)
 
 	// Stop the manager which forces an update of buffered updates.
 	require.NoError(t, mgr.Stop(ctx))
@@ -73,8 +92,8 @@ func TestBufferedUpdates(t *testing.T) {
 			ct.FailNow()
 		}
 
-		assert.EqualValues(ct, 1, interceptor.failed.Load())
-		assert.EqualValues(ct, 2, interceptor.sent.Load())
+		assert.EqualValues(ct, expectedFailure, interceptor.failed.Load())
+		assert.EqualValues(ct, expectedSuccess, interceptor.sent.Load())
 	}, testutil.WaitMedium, testutil.IntervalFast)
 }
 
