enhance retry mechanism with logging and context support

sxwebdev · sxwebdev · commit e35dff8227e2 · 2025-06-05T08:31:46.000+03:00
diff --git a/pkg/db/postgres/logger.go b/pkg/db/postgres/logger.go
@@ -0,0 +1,5 @@
+package postgres
+
+type logger interface {
+	Infof(format string, args ...any)
+}
diff --git a/pkg/db/postgres/postgres.go b/pkg/db/postgres/postgres.go
@@ -44,12 +44,18 @@ func New(ctx context.Context, cfg Config) (*PostgreSQL, error) {
 // CheckConnection checks if the PostgreSQL connection pool is initialized and attempts to ping the database.
 // It retries the ping operation up to 5 times with a 2-second delay between attempts.
 // If the connection pool is not initialized, it returns an error.
-func (p *PostgreSQL) CheckConnection(ctx context.Context) error {
+func (p *PostgreSQL) CheckConnection(ctx context.Context, logger logger) error {
 	if p.DB == nil {
 		return fmt.Errorf("PostgreSQL connection pool is not initialized")
 	}
 
-	return retry.New(retry.WithDelay(time.Second*2), retry.WithMaxAttempts(5)).Do(func() error {
+	return retry.New(
+		retry.WithDelay(time.Second*2),
+		retry.WithMaxAttempts(5),
+		retry.WithPolicy(retry.PolicyLinear),
+		retry.WithLogger(logger),
+		retry.WithContext(ctx),
+	).Do(func() error {
 		return p.DB.Ping(ctx)
 	})
 }
diff --git a/pkg/retry/logger.go b/pkg/retry/logger.go
@@ -0,0 +1,5 @@
+package retry
+
+type logger interface {
+	Infof(format string, args ...any)
+}
diff --git a/pkg/retry/options.go b/pkg/retry/options.go
@@ -7,6 +7,13 @@ import (
 
 type Option func(*Retry)
 
+// WithLogger sets the logger for the Retry instance
+func WithLogger(logger logger) Option {
+	return func(r *Retry) {
+		r.logger = logger
+	}
+}
+
 // WithMaxAttempts sets the maximum number of attempts
 func WithMaxAttempts(maxAttempts int) Option {
 	return func(r *Retry) {
@@ -28,20 +35,33 @@ func WithDelay(delay time.Duration) Option {
 	}
 }
 
-// WithDebug sets the debug mode
-func WithDebug(debug bool) Option {
+// WithContext sets the ctx for Infinite policy retry
+func WithContext(ctx context.Context) Option {
 	return func(r *Retry) {
-		r.debug = debug
+		r.ctx = ctx
 	}
 }
 
-// WithContext sets the ctx for Infinite policy retry
-func WithContext(ctx context.Context) Option {
+// WithOnFailedFn sets the function to be called on failure
+func WithOnFailedFn(fn func()) Option {
 	return func(r *Retry) {
-		r.ctx = ctx
+		r.onFailedFn = fn
 	}
 }
 
+// WithOnSuccessFn sets the function to be called on success
+func WithOnSuccessFn(fn func()) Option {
+	return func(r *Retry) {
+		r.onSuccessFn = fn
+	}
+}
+
+// SetLogger sets the logger for the Retry instance
+func (r *Retry) SetLogger(logger logger) *Retry {
+	r.logger = logger
+	return r
+}
+
 // SetMaxAttempts sets the maximum number of attempts
 func (r *Retry) SetMaxAttempts(maxAttempts int) *Retry {
 	r.maxAttempts = maxAttempts
@@ -60,8 +80,20 @@ func (r *Retry) SetDelay(delay time.Duration) *Retry {
 	return r
 }
 
-// SetDebug sets the debug mode
-func (r *Retry) SetDebug(debug bool) *Retry {
-	r.debug = debug
+// SetContext sets the ctx for Infinite policy retry
+func (r *Retry) SetContext(ctx context.Context) *Retry {
+	r.ctx = ctx
+	return r
+}
+
+// SetOnFailedFn sets the function to be called on failure
+func (r *Retry) SetOnFailedFn(fn func()) *Retry {
+	r.onFailedFn = fn
+	return r
+}
+
+// SetOnSuccessFn sets the function to be called on success
+func (r *Retry) SetOnSuccessFn(fn func()) *Retry {
+	r.onSuccessFn = fn
 	return r
 }
diff --git a/pkg/retry/retry.go b/pkg/retry/retry.go
@@ -10,10 +10,12 @@ import (
 // Retry represents a retry mechanism
 type Retry struct {
 	ctx         context.Context
+	logger      logger
 	maxAttempts int
 	policy      Policy
 	delay       time.Duration
-	debug       bool
+	onFailedFn  func()
+	onSuccessFn func()
 }
 
 // New creates a new Retry instance
@@ -37,16 +39,27 @@ func New(opts ...Option) *Retry {
 
 // Do - performs the retry mechanism
 func (r *Retry) Do(fn func() error) error {
+	if fn == nil {
+		return fmt.Errorf("retry function cannot be nil")
+	}
+
+	var err error
 	switch r.policy {
 	case PolicyLinear:
-		return r.linearRetry(fn)
+		err = r.linearRetry(fn)
 	case PolicyBackoff:
-		return r.backoffRetry(fn)
+		err = r.backoffRetry(fn)
 	case PolicyInfinite:
-		return r.infiniteRetry(fn)
+		err = r.infiniteRetry(fn)
 	default:
-		return fmt.Errorf("unsupported retry policy")
+		err = fmt.Errorf("unsupported retry policy")
+	}
+
+	if err == nil && r.onSuccessFn != nil {
+		r.onSuccessFn()
 	}
+
+	return err
 }
 
 // linearRetry - performs a linear retry mechanism
@@ -57,13 +70,17 @@ func (r *Retry) linearRetry(fn func() error) error {
 			return nil
 		}
 
+		if r.onFailedFn != nil {
+			r.onFailedFn()
+		}
+
 		if errors.Is(err, ErrExit) {
 			return err
 		}
 
 		if attempt < r.maxAttempts {
-			if r.debug {
-				fmt.Printf("linear Retry attempt %d failed, retrying in %s...\n", attempt, r.delay)
+			if r.logger != nil {
+				r.logger.Infof("linear retry attempt %d failed, retrying in %s...", attempt, r.delay)
 			}
 			time.Sleep(r.delay)
 		}
@@ -79,14 +96,18 @@ func (r *Retry) backoffRetry(fn func() error) error {
 			return nil
 		}
 
+		if r.onFailedFn != nil {
+			r.onFailedFn()
+		}
+
 		if errors.Is(err, ErrExit) {
 			return err
 		}
 
 		if attempt < r.maxAttempts {
 			delay := r.delay * (1 << (attempt - 1)) // Увеличение задержки в 2 раза на каждую попытку
-			if r.debug {
-				fmt.Printf("backoff Retry attempt %d failed, retrying in %s...\n", attempt, delay)
+			if r.logger != nil {
+				r.logger.Infof("backoff retry attempt %d failed, retrying in %s...", attempt, delay)
 			}
 			time.Sleep(delay)
 		}
@@ -112,13 +133,17 @@ func (r *Retry) infiniteRetry(fn func() error) error {
 					return
 				}
 
+				if r.onFailedFn != nil {
+					r.onFailedFn()
+				}
+
 				if errors.Is(err, ErrExit) {
 					resCh <- err
 					return
 				}
 
-				if r.debug {
-					fmt.Printf("initnite retry attempt\n")
+				if r.logger != nil {
+					r.logger.Infof("infinite retry attempt failed, retrying in %s...", r.delay)
 				}
 				time.Sleep(r.delay)
 			}
diff --git a/pkg/retry/retry_test.go b/pkg/retry/retry_test.go
@@ -1,55 +1,64 @@
 package retry_test
 
 import (
+	"fmt"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
 	"github.com/tkcrm/modules/pkg/retry"
 )
 
+type logger struct{}
+
+func (l *logger) Infof(format string, args ...any) {
+	fmt.Println(fmt.Sprintf(format, args...))
+}
+
 func TestRetry_Do(t *testing.T) {
+	l := &logger{}
+
 	t.Run("linear retry", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return nil
 		})
 		require.NoError(t, err)
 	})
 	t.Run("backoff retry", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return nil
 		})
 		require.NoError(t, err)
 	})
 
 	t.Run("unsupported retry policy", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.Policy(3)), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.Policy(3)), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return nil
 		})
 		require.Error(t, err)
 	})
 
 	t.Run("linear retry failed", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return retry.ErrRetry
 		})
 		require.Error(t, err)
 	})
 
 	t.Run("backoff retry failed", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return retry.ErrRetry
 		})
 		require.Error(t, err)
 	})
 
 	t.Run("linear retry failed after 3 attempts", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return retry.ErrRetry
 		})
@@ -58,7 +67,7 @@ func TestRetry_Do(t *testing.T) {
 	})
 
 	t.Run("backoff retry failed after 3 attempts", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(3), retry.WithDelay(200*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return retry.ErrRetry
 		})
@@ -67,7 +76,7 @@ func TestRetry_Do(t *testing.T) {
 	})
 
 	t.Run("linear retry failed after 3 attempts with custom delay", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(3), retry.WithDelay(400*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(3), retry.WithDelay(400*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return retry.ErrRetry
 		})
@@ -76,7 +85,7 @@ func TestRetry_Do(t *testing.T) {
 	})
 
 	t.Run("backoff retry failed after 3 attempts with custom delay", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(3), retry.WithDelay(400*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(3), retry.WithDelay(400*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return retry.ErrRetry
 		})
@@ -85,7 +94,7 @@ func TestRetry_Do(t *testing.T) {
 	})
 
 	t.Run("linear retry failed after 3 attempts with custom max attempts", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(5), retry.WithDelay(200*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(5), retry.WithDelay(200*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return retry.ErrRetry
 		})
@@ -94,7 +103,7 @@ func TestRetry_Do(t *testing.T) {
 	})
 
 	t.Run("backoff retry failed after 3 attempts with custom max attempts", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(5), retry.WithDelay(200*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(5), retry.WithDelay(200*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return retry.ErrRetry
 		})
@@ -103,7 +112,7 @@ func TestRetry_Do(t *testing.T) {
 	})
 
 	t.Run("linear retry failed after 3 attempts with custom max attempts and delay", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(5), retry.WithDelay(400*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(5), retry.WithDelay(400*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return retry.ErrRetry
 		})
@@ -112,7 +121,7 @@ func TestRetry_Do(t *testing.T) {
 	})
 
 	t.Run("backoff retry failed after 3 attempts with custom max attempts and delay", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(5), retry.WithDelay(400*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(5), retry.WithDelay(400*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return retry.ErrRetry
 		})
@@ -121,7 +130,7 @@ func TestRetry_Do(t *testing.T) {
 	})
 
 	t.Run("linear retry failed after 3 attempts with custom max attempts and delay and custom retry function", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(5), retry.WithDelay(400*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithMaxAttempts(5), retry.WithDelay(400*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return retry.ErrRetry
 		})
@@ -130,7 +139,7 @@ func TestRetry_Do(t *testing.T) {
 	})
 
 	t.Run("backoff retry failed after 3 attempts with custom max attempts and delay and custom retry function", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(5), retry.WithDelay(400*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyBackoff), retry.WithMaxAttempts(5), retry.WithDelay(400*time.Millisecond)).SetLogger(l)
 		err := r.Do(func() error {
 			return retry.ErrRetry
 		})
@@ -139,7 +148,7 @@ func TestRetry_Do(t *testing.T) {
 	})
 
 	t.Run("linear retry failed with exit error", func(t *testing.T) {
-		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithDelay(200*time.Millisecond)).SetDebug(true)
+		r := retry.New(retry.WithPolicy(retry.PolicyLinear), retry.WithDelay(200*time.Millisecond)).SetLogger(l)
 		var attempt int
 		err := r.Do(func() error {
 			attempt++
