fix: improve http connection pooling for smtp notifications

kacpersaw · kacpersaw · commit c09999a6e9b5 · 2025-10-31T10:44:57.000Z
diff --git a/cli/exp_scaletest_notifications.go b/cli/exp_scaletest_notifications.go
@@ -142,6 +142,16 @@ func (r *RootCmd) scaletestNotifications() *serpent.Command {
 				triggerTimes[id] = make(chan time.Time, 1)
 			}
 
+			smtpHTTPTransport := &http.Transport{
+				MaxConnsPerHost:     512,
+				MaxIdleConnsPerHost: 512,
+				MaxIdleConns:        2048,
+				IdleConnTimeout:     60 * time.Second,
+			}
+			smtpHTTPClient := &http.Client{
+				Transport: smtpHTTPTransport,
+			}
+
 			configs := make([]notifications.Config, 0, userCount)
 			for range templateAdminCount {
 				config := notifications.Config{
@@ -157,6 +167,7 @@ func (r *RootCmd) scaletestNotifications() *serpent.Command {
 					Metrics:                  metrics,
 					SMTPApiURL:               smtpAPIURL,
 					SMTPRequestTimeout:       smtpRequestTimeout,
+					SMTPHttpClient:           smtpHTTPClient,
 				}
 				if err := config.Validate(); err != nil {
 					return xerrors.Errorf("validate config: %w", err)
diff --git a/scaletest/notifications/config.go b/scaletest/notifications/config.go
@@ -1,6 +1,7 @@
 package notifications
 
 import (
+	"net/http"
 	"sync"
 	"time"
 
@@ -40,6 +41,9 @@ type Config struct {
 
 	// SMTPRequestTimeout is the timeout for SMTP requests.
 	SMTPRequestTimeout time.Duration `json:"smtp_request_timeout"`
+
+	// SMTPHttpClient is the HTTP client for SMTP requests.
+	SMTPHttpClient *http.Client `json:"-"`
 }
 
 func (c Config) Validate() error {
@@ -68,6 +72,10 @@ func (c Config) Validate() error {
 		return xerrors.New("smtp_request_timeout must be set if smtp_api_url is set")
 	}
 
+	if c.SMTPApiURL != "" && c.SMTPHttpClient == nil {
+		return xerrors.New("smtp_http_client must be set if smtp_api_url is set")
+	}
+
 	if c.DialTimeout <= 0 {
 		return xerrors.New("dial_timeout must be greater than 0")
 	}
diff --git a/scaletest/notifications/run.go b/scaletest/notifications/run.go
@@ -298,15 +298,16 @@ func (r *Runner) watchNotificationsSMTP(ctx context.Context, user codersdk.User,
 	receivedNotifications := make(map[uuid.UUID]struct{})
 
 	apiURL := fmt.Sprintf("%s/messages?email=%s", r.cfg.SMTPApiURL, user.Email)
-	httpClient := &http.Client{
-		Timeout: r.cfg.SMTPRequestTimeout,
-	}
+	httpClient := r.cfg.SMTPHttpClient
 
 	const smtpPollInterval = 2 * time.Second
 	done := xerrors.New("done")
 
 	tkr := r.clock.TickerFunc(ctx, smtpPollInterval, func() error {
-		req, err := http.NewRequestWithContext(ctx, http.MethodGet, apiURL, nil)
+		reqCtx, cancel := context.WithTimeout(ctx, r.cfg.SMTPRequestTimeout)
+		defer cancel()
+
+		req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, apiURL, nil)
 		if err != nil {
 			logger.Error(ctx, "create SMTP API request", slog.Error(err))
 			r.cfg.Metrics.AddError("smtp_create_request")
@@ -317,14 +318,16 @@ func (r *Runner) watchNotificationsSMTP(ctx context.Context, user codersdk.User,
 		if err != nil {
 			logger.Error(ctx, "poll smtp api for notifications", slog.Error(err))
 			r.cfg.Metrics.AddError("smtp_poll")
-			return xerrors.Errorf("poll smtp api: %w", err)
+			return nil
 		}
 
 		if resp.StatusCode != http.StatusOK {
+			// discard the response to allow reusing of the connection
+			_, _ = io.Copy(io.Discard, resp.Body)
 			_ = resp.Body.Close()
 			logger.Error(ctx, "smtp api returned non-200 status", slog.F("status", resp.StatusCode))
 			r.cfg.Metrics.AddError("smtp_bad_status")
-			return xerrors.Errorf("smtp api returned status %d", resp.StatusCode)
+			return nil
 		}
 
 		var summaries []smtpmock.EmailSummary
