Jul 8, 2024 · Jun 11, 2024 · Jun 11, 2024 · Jun 11, 2024 · Jun 11, 2024 · Jun 12, 2024
diff --git a/coderd/notifications/manager.go b/coderd/notifications/manager.go
 notifier *notifier
 handlers map[database.NotificationMethod]Handler

 success, failure chan dispatchResult

 runOnce  sync.Once
 stopOnce sync.Once
 stop     chan any
 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),

 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.
 // 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)
 }
 }
 })
 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 {
 // 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
 "encoding/json"
 "sync/atomic"
 "testing"
 "time"

 "github.com/coder/serpent"
 "github.com/google/uuid"
 "github.com/stretchr/testify/assert"
 "github.com/stretchr/testify/require"
 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{
 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))
 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)
 }
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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
Expand All		@@ -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),

Expand DownExpand Up		@@ -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.
Expand All		@@ -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)
		}
		}
		})
Expand All		@@ -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 {
Expand All		@@ -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 {
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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"
Expand All		@@ -31,11 +33,14 @@ func TestBufferedUpdates(t *testing.T) {
		if !dbtestutil.WillUsePostgres() {
		t.Skip("This test requires postgres")
spikecurtis marked this conversation as resolved. Show resolvedHide resolved
		}

		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{
Expand All		@@ -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))
Copy link Contributor spikecurtisJul 3, 2024 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. there is a race condition here, which is that above we wait for notifications to be sent to santa, but sending to santa happensbefore sending on the`success` and`failure` channels. So, triggering stop and the final bulk update might miss one or more success/failure updates and cause this test to flake. Quartz can help here. The manager's bulk update and the notifier's processing loop depend on Tickers. If you convert them to use a`quartz.TickerFunc`, then you can control the mock clock and trigger the processing loop and bulk update, and wait for them to finish. This lets us get rid of the`Eventually` that checks for dispatched notifications (Quartz allows you to wait until the tick processing is done, so there is no race in just asserting santa got what we expected). It also means we can get rid of the`Eventually` below for the same reason --- we can wait for bulk processing to complete before asserting the`interceptor`'s state. Copy link ContributorAuthor dannykoppingJul 3, 2024 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. I've got a follow-up item to implement Quartz in a subsequent PR. Copy link ContributorAuthor dannykoppingJul 4, 2024 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. Addressed in`613e074`
Expand All		@@ -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)
		}

Expand Down