diff --git a/scaletest/notifications/config.go b/scaletest/notifications/config.go new file mode 100644 index 0000000000..ac1c6da49a --- /dev/null +++ b/scaletest/notifications/config.go @@ -0,0 +1,71 @@ +package notifications + +import ( + "sync" + "time" + + "golang.org/x/xerrors" + + "github.com/google/uuid" + + "github.com/coder/coder/v2/scaletest/createusers" +) + +type Config struct { + // User is the configuration for the user to create. + User createusers.Config `json:"user"` + + // Roles are the roles to assign to the user. + Roles []string `json:"roles"` + + // NotificationTimeout is how long to wait for notifications after triggering. + NotificationTimeout time.Duration `json:"notification_timeout"` + + // DialTimeout is how long to wait for websocket connection. + DialTimeout time.Duration `json:"dial_timeout"` + + // ExpectedNotifications maps notification template IDs to channels + // that receive the trigger time for each notification. + ExpectedNotifications map[uuid.UUID]chan time.Time `json:"-"` + + Metrics *Metrics `json:"-"` + + // DialBarrier ensures all runners are connected before notifications are triggered. + DialBarrier *sync.WaitGroup `json:"-"` + + // ReceivingWatchBarrier is the barrier for receiving users. Regular users wait on this to disconnect after receiving users complete. + ReceivingWatchBarrier *sync.WaitGroup `json:"-"` +} + +func (c Config) Validate() error { + // The runner always needs an org; ensure we propagate it into the user config. + if c.User.OrganizationID == uuid.Nil { + return xerrors.New("user organization_id must be set") + } + + if err := c.User.Validate(); err != nil { + return xerrors.Errorf("user config: %w", err) + } + + if c.DialBarrier == nil { + return xerrors.New("dial barrier must be set") + } + + if c.ReceivingWatchBarrier == nil { + return xerrors.New("receiving_watch_barrier must be set") + } + + if c.NotificationTimeout <= 0 { + return xerrors.New("notification_timeout must be greater than 0") + } + + if c.DialTimeout <= 0 { + return xerrors.New("dial_timeout must be greater than 0") + } + + if c.Metrics == nil { + return xerrors.New("metrics must be set") + } + + return nil +} diff --git a/scaletest/notifications/metrics.go b/scaletest/notifications/metrics.go new file mode 100644 index 0000000000..c9e7374250 --- /dev/null +++ b/scaletest/notifications/metrics.go @@ -0,0 +1,58 @@ +package notifications + +import ( + "time" + + "github.com/prometheus/client_golang/prometheus" +) + +type Metrics struct { + notificationLatency *prometheus.HistogramVec + notificationErrors *prometheus.CounterVec + missedNotifications *prometheus.CounterVec +} + +func NewMetrics(reg prometheus.Registerer) *Metrics { + if reg == nil { + reg = prometheus.DefaultRegisterer + } + + latency := prometheus.NewHistogramVec(prometheus.HistogramOpts{ + Namespace: "coderd", + Subsystem: "scaletest", + Name: "notification_delivery_latency_seconds", + Help: "Time between notification-creating action and receipt of notification by client", + }, []string{"username", "notification_type"}) + errors := prometheus.NewCounterVec(prometheus.CounterOpts{ + Namespace: "coderd", + Subsystem: "scaletest", + Name: "notification_delivery_errors_total", + Help: "Total number of notification delivery errors", + }, []string{"username", "action"}) + missed := prometheus.NewCounterVec(prometheus.CounterOpts{ + Namespace: "coderd", + Subsystem: "scaletest", + Name: "notification_delivery_missed_total", + Help: "Total number of missed notifications", + }, []string{"username"}) + + reg.MustRegister(latency, errors, missed) + + return &Metrics{ + notificationLatency: latency, + notificationErrors: errors, + missedNotifications: missed, + } +} + +func (m *Metrics) RecordLatency(latency time.Duration, username, notificationType string) { + m.notificationLatency.WithLabelValues(username, notificationType).Observe(latency.Seconds()) +} + +func (m *Metrics) AddError(username, action string) { + m.notificationErrors.WithLabelValues(username, action).Inc() +} + +func (m *Metrics) RecordMissed(username string) { + m.missedNotifications.WithLabelValues(username).Inc() +} diff --git a/scaletest/notifications/run.go b/scaletest/notifications/run.go new file mode 100644 index 0000000000..d3d68e78ac --- /dev/null +++ b/scaletest/notifications/run.go @@ -0,0 +1,263 @@ +package notifications + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "time" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "cdr.dev/slog" + "cdr.dev/slog/sloggers/sloghuman" + + "github.com/coder/coder/v2/coderd/tracing" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/scaletest/createusers" + "github.com/coder/coder/v2/scaletest/harness" + "github.com/coder/coder/v2/scaletest/loadtestutil" + "github.com/coder/websocket" +) + +type Runner struct { + client *codersdk.Client + cfg Config + + createUserRunner *createusers.Runner + + // notificationLatencies stores the latency for each notification type + notificationLatencies map[uuid.UUID]time.Duration +} + +func NewRunner(client *codersdk.Client, cfg Config) *Runner { + return &Runner{ + client: client, + cfg: cfg, + notificationLatencies: make(map[uuid.UUID]time.Duration), + } +} + +var ( + _ harness.Runnable = &Runner{} + _ harness.Cleanable = &Runner{} + _ harness.Collectable = &Runner{} +) + +func (r *Runner) Run(ctx context.Context, id string, logs io.Writer) error { + ctx, span := tracing.StartSpan(ctx) + defer span.End() + + reachedBarrier := false + defer func() { + if !reachedBarrier { + r.cfg.DialBarrier.Done() + } + }() + + reachedReceivingWatchBarrier := false + defer func() { + if len(r.cfg.ExpectedNotifications) > 0 && !reachedReceivingWatchBarrier { + r.cfg.ReceivingWatchBarrier.Done() + } + }() + + logs = loadtestutil.NewSyncWriter(logs) + logger := slog.Make(sloghuman.Sink(logs)).Leveled(slog.LevelDebug) + r.client.SetLogger(logger) + r.client.SetLogBodies(true) + + r.createUserRunner = createusers.NewRunner(r.client, r.cfg.User) + newUserAndToken, err := r.createUserRunner.RunReturningUser(ctx, id, logs) + if err != nil { + r.cfg.Metrics.AddError("", "create_user") + return xerrors.Errorf("create user: %w", err) + } + newUser := newUserAndToken.User + newUserClient := codersdk.New(r.client.URL, + codersdk.WithSessionToken(newUserAndToken.SessionToken), + codersdk.WithLogger(logger), + codersdk.WithLogBodies()) + + logger.Info(ctx, "runner user created", slog.F("username", newUser.Username), slog.F("user_id", newUser.ID.String())) + + if len(r.cfg.Roles) > 0 { + logger.Info(ctx, "assigning roles to user", slog.F("roles", r.cfg.Roles)) + + _, err := r.client.UpdateUserRoles(ctx, newUser.ID.String(), codersdk.UpdateRoles{ + Roles: r.cfg.Roles, + }) + if err != nil { + r.cfg.Metrics.AddError(newUser.Username, "assign_roles") + return xerrors.Errorf("assign roles: %w", err) + } + } + + logger.Info(ctx, "notification runner is ready") + + dialCtx, cancel := context.WithTimeout(ctx, r.cfg.DialTimeout) + defer cancel() + + logger.Info(ctx, "connecting to notification websocket") + conn, err := r.dialNotificationWebsocket(dialCtx, newUserClient, newUser, logger) + if err != nil { + return xerrors.Errorf("dial notification websocket: %w", err) + } + defer conn.Close(websocket.StatusNormalClosure, "done") + logger.Info(ctx, "connected to notification websocket") + + reachedBarrier = true + r.cfg.DialBarrier.Done() + r.cfg.DialBarrier.Wait() + + if len(r.cfg.ExpectedNotifications) == 0 { + logger.Info(ctx, "maintaining websocket connection, waiting for receiving users to complete") + + // Wait for receiving users to complete + done := make(chan struct{}) + go func() { + r.cfg.ReceivingWatchBarrier.Wait() + close(done) + }() + + select { + case <-done: + logger.Info(ctx, "receiving users complete, closing connection") + case <-ctx.Done(): + logger.Info(ctx, "context canceled, closing connection") + } + return nil + } + + logger.Info(ctx, "waiting for notifications", slog.F("timeout", r.cfg.NotificationTimeout)) + + watchCtx, cancel := context.WithTimeout(ctx, r.cfg.NotificationTimeout) + defer cancel() + + if err := r.watchNotifications(watchCtx, conn, newUser, logger, r.cfg.ExpectedNotifications); err != nil { + return xerrors.Errorf("notification watch failed: %w", err) + } + + reachedReceivingWatchBarrier = true + r.cfg.ReceivingWatchBarrier.Done() + + return nil +} + +func (r *Runner) Cleanup(ctx context.Context, id string, logs io.Writer) error { + if r.createUserRunner != nil { + _, _ = fmt.Fprintln(logs, "Cleaning up user...") + if err := r.createUserRunner.Cleanup(ctx, id, logs); err != nil { + return xerrors.Errorf("cleanup user: %w", err) + } + } + + return nil +} + +const NotificationDeliveryLatencyMetric = "notification_delivery_latency_seconds" + +func (r *Runner) GetMetrics() map[string]any { + return map[string]any{ + NotificationDeliveryLatencyMetric: r.notificationLatencies, + } +} + +func (r *Runner) dialNotificationWebsocket(ctx context.Context, client *codersdk.Client, user codersdk.User, logger slog.Logger) (*websocket.Conn, error) { + u, err := client.URL.Parse("/api/v2/notifications/inbox/watch") + if err != nil { + logger.Error(ctx, "parse notification URL", slog.Error(err)) + r.cfg.Metrics.AddError(user.Username, "parse_url") + return nil, xerrors.Errorf("parse notification URL: %w", err) + } + + conn, resp, err := websocket.Dial(ctx, u.String(), &websocket.DialOptions{ + HTTPHeader: http.Header{ + "Coder-Session-Token": []string{client.SessionToken()}, + }, + }) + if err != nil { + if resp != nil { + defer resp.Body.Close() + if resp.StatusCode != http.StatusSwitchingProtocols { + err = codersdk.ReadBodyAsError(resp) + } + } + logger.Error(ctx, "dial notification websocket", slog.Error(err)) + r.cfg.Metrics.AddError(user.Username, "dial") + return nil, xerrors.Errorf("dial notification websocket: %w", err) + } + + return conn, nil +} + +// watchNotifications reads notifications from the websocket and returns error or nil +// once all expected notifications are received. +func (r *Runner) watchNotifications(ctx context.Context, conn *websocket.Conn, user codersdk.User, logger slog.Logger, expectedNotifications map[uuid.UUID]chan time.Time) error { + logger.Info(ctx, "waiting for notifications", + slog.F("username", user.Username), + slog.F("expected_count", len(expectedNotifications))) + + receivedNotifications := make(map[uuid.UUID]struct{}) + + for { + select { + case <-ctx.Done(): + return xerrors.Errorf("context canceled while waiting for notifications: %w", ctx.Err()) + default: + } + + if len(receivedNotifications) == len(expectedNotifications) { + logger.Info(ctx, "received all expected notifications") + return nil + } + + notif, err := readNotification(ctx, conn) + if err != nil { + logger.Error(ctx, "read notification", slog.Error(err)) + r.cfg.Metrics.AddError(user.Username, "read_notification") + return xerrors.Errorf("read notification: %w", err) + } + + templateID := notif.Notification.TemplateID + if triggerTimeChan, exists := expectedNotifications[templateID]; exists { + if _, exists := receivedNotifications[templateID]; !exists { + receiptTime := time.Now() + select { + case triggerTime := <-triggerTimeChan: + latency := receiptTime.Sub(triggerTime) + r.notificationLatencies[templateID] = latency + r.cfg.Metrics.RecordLatency(latency, user.Username, templateID.String()) + receivedNotifications[templateID] = struct{}{} + + logger.Info(ctx, "received expected notification", + slog.F("template_id", templateID), + slog.F("title", notif.Notification.Title), + slog.F("latency", latency)) + case <-ctx.Done(): + return xerrors.Errorf("context canceled while waiting for trigger time: %w", ctx.Err()) + } + } + } else { + logger.Debug(ctx, "received notification not being tested", + slog.F("template_id", templateID), + slog.F("title", notif.Notification.Title)) + } + } +} + +func readNotification(ctx context.Context, conn *websocket.Conn) (codersdk.GetInboxNotificationResponse, error) { + _, message, err := conn.Read(ctx) + if err != nil { + return codersdk.GetInboxNotificationResponse{}, err + } + + var notif codersdk.GetInboxNotificationResponse + if err := json.Unmarshal(message, ¬if); err != nil { + return codersdk.GetInboxNotificationResponse{}, xerrors.Errorf("unmarshal notification: %w", err) + } + + return notif, nil +} diff --git a/scaletest/notifications/run_test.go b/scaletest/notifications/run_test.go new file mode 100644 index 0000000000..e94e6d82ea --- /dev/null +++ b/scaletest/notifications/run_test.go @@ -0,0 +1,221 @@ +package notifications_test + +import ( + "io" + "strconv" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/sync/errgroup" + "golang.org/x/xerrors" + + "github.com/coder/serpent" + + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + notificationsLib "github.com/coder/coder/v2/coderd/notifications" + "github.com/coder/coder/v2/coderd/notifications/dispatch" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/scaletest/createusers" + "github.com/coder/coder/v2/scaletest/notifications" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +func TestRun(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + logger := testutil.Logger(t) + db, ps := dbtestutil.NewDB(t) + + // Setup notifications manager with inbox handler + cfg := defaultNotificationsConfig(database.NotificationMethodSmtp) + mgr, err := notificationsLib.NewManager( + cfg, + db, + ps, + defaultHelpers(), + notificationsLib.NewMetrics(prometheus.NewRegistry()), + logger.Named("manager"), + ) + require.NoError(t, err) + + mgr.WithHandlers(map[database.NotificationMethod]notificationsLib.Handler{ + database.NotificationMethodInbox: dispatch.NewInboxHandler(logger.Named("inbox"), db, ps), + }) + t.Cleanup(func() { + assert.NoError(t, mgr.Stop(dbauthz.AsNotifier(ctx))) + }) + mgr.Run(dbauthz.AsNotifier(ctx)) + + enqueuer, err := notificationsLib.NewStoreEnqueuer( + cfg, + db, + defaultHelpers(), + logger.Named("enqueuer"), + quartz.NewReal(), + ) + require.NoError(t, err) + + client := coderdtest.New(t, &coderdtest.Options{ + Database: db, + Pubsub: ps, + NotificationsEnqueuer: enqueuer, + }) + firstUser := coderdtest.CreateFirstUser(t, client) + + const numReceivingUsers = 2 + const numRegularUsers = 2 + dialBarrier := new(sync.WaitGroup) + receivingWatchBarrier := new(sync.WaitGroup) + dialBarrier.Add(numReceivingUsers + numRegularUsers) + receivingWatchBarrier.Add(numReceivingUsers) + metrics := notifications.NewMetrics(prometheus.NewRegistry()) + + eg, runCtx := errgroup.WithContext(ctx) + + expectedNotifications := map[uuid.UUID]chan time.Time{ + notificationsLib.TemplateUserAccountCreated: make(chan time.Time, 1), + notificationsLib.TemplateUserAccountDeleted: make(chan time.Time, 1), + } + + // Start receiving runners who will receive notifications + receivingRunners := make([]*notifications.Runner, 0, numReceivingUsers) + for i := range numReceivingUsers { + runnerCfg := notifications.Config{ + User: createusers.Config{ + OrganizationID: firstUser.OrganizationID, + }, + Roles: []string{codersdk.RoleOwner}, + NotificationTimeout: testutil.WaitLong, + DialTimeout: testutil.WaitLong, + Metrics: metrics, + DialBarrier: dialBarrier, + ReceivingWatchBarrier: receivingWatchBarrier, + ExpectedNotifications: expectedNotifications, + } + err := runnerCfg.Validate() + require.NoError(t, err) + + runner := notifications.NewRunner(client, runnerCfg) + receivingRunners = append(receivingRunners, runner) + eg.Go(func() error { + return runner.Run(runCtx, "receiving-"+strconv.Itoa(i), io.Discard) + }) + } + + // Start regular user runners who will maintain websocket connections + regularRunners := make([]*notifications.Runner, 0, numRegularUsers) + for i := range numRegularUsers { + runnerCfg := notifications.Config{ + User: createusers.Config{ + OrganizationID: firstUser.OrganizationID, + }, + Roles: []string{}, + NotificationTimeout: testutil.WaitLong, + DialTimeout: testutil.WaitLong, + Metrics: metrics, + DialBarrier: dialBarrier, + ReceivingWatchBarrier: receivingWatchBarrier, + } + err := runnerCfg.Validate() + require.NoError(t, err) + + runner := notifications.NewRunner(client, runnerCfg) + regularRunners = append(regularRunners, runner) + eg.Go(func() error { + return runner.Run(runCtx, "regular-"+strconv.Itoa(i), io.Discard) + }) + } + + // Trigger notifications by creating and deleting a user + eg.Go(func() error { + // Wait for all runners to connect + dialBarrier.Wait() + + createTime := time.Now() + newUser, err := client.CreateUserWithOrgs(runCtx, codersdk.CreateUserRequestWithOrgs{ + OrganizationIDs: []uuid.UUID{firstUser.OrganizationID}, + Email: "test-user@coder.com", + Username: "test-user", + Password: "SomeSecurePassword!", + }) + if err != nil { + return xerrors.Errorf("create test user: %w", err) + } + expectedNotifications[notificationsLib.TemplateUserAccountCreated] <- createTime + + deleteTime := time.Now() + if err := client.DeleteUser(runCtx, newUser.ID); err != nil { + return xerrors.Errorf("delete test user: %w", err) + } + expectedNotifications[notificationsLib.TemplateUserAccountDeleted] <- deleteTime + + close(expectedNotifications[notificationsLib.TemplateUserAccountCreated]) + close(expectedNotifications[notificationsLib.TemplateUserAccountDeleted]) + + return nil + }) + + err = eg.Wait() + require.NoError(t, err, "runner execution should complete successfully") + + cleanupEg, cleanupCtx := errgroup.WithContext(ctx) + for i, runner := range receivingRunners { + cleanupEg.Go(func() error { + return runner.Cleanup(cleanupCtx, "receiving-"+strconv.Itoa(i), io.Discard) + }) + } + for i, runner := range regularRunners { + cleanupEg.Go(func() error { + return runner.Cleanup(cleanupCtx, "regular-"+strconv.Itoa(i), io.Discard) + }) + } + err = cleanupEg.Wait() + require.NoError(t, err) + + users, err := client.Users(ctx, codersdk.UsersRequest{}) + require.NoError(t, err) + require.Len(t, users.Users, 1) + require.Equal(t, firstUser.UserID, users.Users[0].ID) + + for _, runner := range receivingRunners { + runnerMetrics := runner.GetMetrics()[notifications.NotificationDeliveryLatencyMetric].(map[uuid.UUID]time.Duration) + require.Contains(t, runnerMetrics, notificationsLib.TemplateUserAccountCreated) + require.Contains(t, runnerMetrics, notificationsLib.TemplateUserAccountDeleted) + } +} + +func defaultNotificationsConfig(method database.NotificationMethod) codersdk.NotificationsConfig { + return codersdk.NotificationsConfig{ + Method: serpent.String(method), + MaxSendAttempts: 5, + FetchInterval: serpent.Duration(time.Millisecond * 100), + StoreSyncInterval: serpent.Duration(time.Millisecond * 200), + LeasePeriod: serpent.Duration(time.Second * 10), + DispatchTimeout: serpent.Duration(time.Second * 5), + RetryInterval: serpent.Duration(time.Millisecond * 50), + LeaseCount: 10, + StoreSyncBufferSize: 50, + Inbox: codersdk.NotificationsInboxConfig{ + Enabled: serpent.Bool(true), + }, + } +} + +func defaultHelpers() map[string]any { + return map[string]any{ + "base_url": func() string { return "http://test.com" }, + "current_year": func() string { return "2024" }, + "logo_url": func() string { return "https://coder.com/coder-logo-horizontal.png" }, + "app_name": func() string { return "Coder" }, + } +}