changes as per comments, add tests

make base delay constant 1 second,
decrement retryAttempt in notificationListenerLoop,
hardcode sleep duration in shutdown,
slight change in logic - cap delays to max of 120 seconds, but allow for jitter in that too,
add tests according to calculations

Signed-off-by: pranshu-raj-211 <[email protected]>

Author: pranshu-raj-211

dbos/system_database.go

@@ -136,10 +136,10 @@ const (
 	_DBOS_WORKFLOW_EVENTS_CHANNEL = "dbos_workflow_events_channel"
 
 	// Database retry timeouts
-	_DB_CONNECTION_RETRY_BASE_DELAY  = 500 * time.Millisecond
+	_DB_CONNECTION_RETRY_BASE_DELAY  = 1 * time.Second
 	_DB_CONNECTION_RETRY_FACTOR      = 2
-	_DB_CONNECTION_RETRY_MAX_RETRIES = 3
-	_DB_CONNECTION_MAX_DELAY         = 10000 * time.Millisecond
+	_DB_CONNECTION_RETRY_MAX_RETRIES = 10
+	_DB_CONNECTION_MAX_DELAY         = 120 * time.Second
 	_DB_RETRY_INTERVAL               = 1 * time.Second
 )
 
@@ -293,7 +293,7 @@ func (s *sysDB) shutdown(ctx context.Context, timeout time.Duration) {
 	// Allow pgx health checks to complete
 	// https://github.com/jackc/pgx/blob/15bca4a4e14e0049777c1245dba4c16300fe4fd0/pgxpool/pool.go#L417
 	// These trigger go-leak alerts
-	time.Sleep(backoffWithJitter(0))
+	time.Sleep(500 * time.Second)
 
 	s.launched = false
 }
@@ -1532,6 +1532,10 @@ func (s *sysDB) notificationListenerLoop(ctx context.Context) {
 			time.Sleep(backoffWithJitter(retryAttempt))
 			retryAttempt += 1
 			continue
+		} else {
+			if retryAttempt > 0 {
+				retryAttempt -= 1
+			}
 		}
 	}
 }
@@ -2337,14 +2341,11 @@ func (qb *queryBuilder) addWhereLessEqual(column string, value any) {
 }
 
 func backoffWithJitter(retryAttempt int) time.Duration {
+	exp := float64(_DB_CONNECTION_RETRY_BASE_DELAY) * math.Pow(_DB_CONNECTION_RETRY_FACTOR, float64(retryAttempt))
 	// cap backoff to max number of retries, then do a fixed time delay
 	// expected retryAttempt to initially be 0, so >= used
-	if retryAttempt >= _DB_CONNECTION_RETRY_MAX_RETRIES {
-		return _DB_CONNECTION_RETRY_BASE_DELAY
-	}
-	exp := float64(_DB_CONNECTION_RETRY_BASE_DELAY) * math.Pow(_DB_CONNECTION_RETRY_FACTOR, float64(retryAttempt))
 	// cap delay to maximum of _DB_CONNECTION_MAX_DELAY milliseconds
-	if exp > float64(_DB_CONNECTION_MAX_DELAY) {
+	if retryAttempt >= _DB_CONNECTION_RETRY_MAX_RETRIES || exp > float64(_DB_CONNECTION_MAX_DELAY) {
 		exp = float64(_DB_CONNECTION_MAX_DELAY)
 	}
 

dbos/system_database_test.go

@@ -0,0 +1,51 @@
+package dbos
+
+import (
+	"testing"
+	"time"
+)
+
+var backoffWithJitterTestcases = []struct {
+	name         string
+	retryAttempt int
+	wantMin      time.Duration
+	wantMax      time.Duration
+}{
+	{
+		name:         "first retry attempt (0)",
+		retryAttempt: 0,
+		wantMin:      750 * time.Millisecond,
+		wantMax:      1250 * time.Millisecond,
+	},
+	{
+		name:         "second retry attempt (1)",
+		retryAttempt: 1,
+		wantMin:      1500 * time.Millisecond,
+		wantMax:      2500 * time.Millisecond,
+	},
+	{
+		name:         "ninth retry attempt (8)",
+		retryAttempt: 8,
+		wantMin:      90 * time.Second,
+		wantMax:      150 * time.Second,
+	},
+	{
+		name:         "exceeds max retries",
+		retryAttempt: 10,
+		wantMin:      90 * time.Second,
+		wantMax:      150 * time.Second,
+	},
+}
+
+func TestBackoffWithJitter(t *testing.T) {
+	for _, testcase := range backoffWithJitterTestcases {
+		t.Run(testcase.name, func(t *testing.T) {
+			got := backoffWithJitter(testcase.retryAttempt)
+
+			if got < testcase.wantMin || got > testcase.wantMax {
+				t.Errorf("Should be between %v and %v, got=%v, attempt=%v",
+					testcase.wantMin, testcase.wantMax, got, testcase.retryAttempt)
+			}
+		})
+	}
+}