Remove underlay type tracking in dial policy

Author: plorenz

dial_policy.go

@@ -25,137 +25,109 @@ type BackoffConfig struct {
 	MinStableDuration time.Duration
 }
 
+// DefaultBackoffConfig provides sensible defaults for BackoffConfig.
 var DefaultBackoffConfig = BackoffConfig{
 	BaseDelay:         2 * time.Second,
 	MaxDelay:          time.Minute,
 	MinStableDuration: 5 * time.Second,
 }
 
-// dialState tracks dial state for a specific underlay type.
-type dialState struct {
-	consecutiveFailures uint32
-	lastSuccess         time.Time
-}
-
 // BackoffDialPolicy wraps a DialUnderlayFactory with exponential backoff retry logic.
-// It tracks consecutive failures per underlay type and detects short-lived connections.
+// It tracks consecutive failures and detects short-lived connections.
 type BackoffDialPolicy struct {
 	Dialer  DialUnderlayFactory
 	Backoff BackoffConfig
 	mu      sync.Mutex
-	state   map[string]*dialState
+	consecutiveFailures uint32
+	lastSuccess         time.Time
 }
 
+// NewBackoffDialPolicy creates a BackoffDialPolicy with default configuration.
 func NewBackoffDialPolicy(dialer DialUnderlayFactory) *BackoffDialPolicy {
 	return &BackoffDialPolicy{
 		Dialer:  dialer,
 		Backoff: DefaultBackoffConfig,
-		state:   map[string]*dialState{},
 	}
 }
 
+// NewBackoffDialPolicyWithConfig creates a BackoffDialPolicy with the given configuration.
 func NewBackoffDialPolicyWithConfig(dialer DialUnderlayFactory, config BackoffConfig) *BackoffDialPolicy {
 	return &BackoffDialPolicy{
 		Dialer:  dialer,
 		Backoff: config,
-		state:   map[string]*dialState{},
 	}
 }
 
 // recordShortLivedConnection checks if the last successful connection was short-lived
 // (died before MinStableDuration elapsed) and treats it as a failure for backoff purposes.
 // If the connection lived long enough to be considered stable, the failure count is reset.
-func (self *BackoffDialPolicy) recordShortLivedConnection(underlayType string) {
+func (self *BackoffDialPolicy) recordShortLivedConnection() {
 	self.mu.Lock()
 	defer self.mu.Unlock()
 
-	s, ok := self.state[underlayType]
-	if !ok {
-		return
-	}
-
-	if s.lastSuccess.IsZero() {
+	if self.lastSuccess.IsZero() {
 		return
 	}
 
-	if time.Since(s.lastSuccess) < self.Backoff.MinStableDuration {
-		s.consecutiveFailures++
+	if time.Since(self.lastSuccess) < self.Backoff.MinStableDuration {
+		self.consecutiveFailures++
 	} else {
-		s.consecutiveFailures = 0
+		self.consecutiveFailures = 0
 	}
-	s.lastSuccess = time.Time{}
+	self.lastSuccess = time.Time{}
 }
 
-func (self *BackoffDialPolicy) getBackoffDelay(underlayType string) time.Duration {
+func (self *BackoffDialPolicy) getBackoffDelay() time.Duration {
 	self.mu.Lock()
 	defer self.mu.Unlock()
 
-	s, ok := self.state[underlayType]
-	if !ok {
-		return 0
-	}
-
-	if s.consecutiveFailures == 0 {
+	if self.consecutiveFailures == 0 {
 		return 0
 	}
 
 	// Exponential backoff: baseDelay * 2^(failures-1), capped at maxDelay
-	delay := self.Backoff.BaseDelay * (1 << min(s.consecutiveFailures-1, 30))
+	delay := self.Backoff.BaseDelay * (1 << min(self.consecutiveFailures-1, 30))
 	return min(delay, self.Backoff.MaxDelay)
 }
 
-func (self *BackoffDialPolicy) recordSuccess(underlayType string) {
+func (self *BackoffDialPolicy) recordSuccess() {
 	self.mu.Lock()
 	defer self.mu.Unlock()
-	s, ok := self.state[underlayType]
-	if !ok {
-		s = &dialState{}
-		self.state[underlayType] = s
-	}
-	s.lastSuccess = time.Now()
+	self.lastSuccess = time.Now()
 	// Don't reset consecutiveFailures yet — wait to see if connection is stable.
-	// Checked on next call to getBackoffDelay.
+	// Checked on next call to recordShortLivedConnection.
 }
 
-func (self *BackoffDialPolicy) recordFailure(underlayType string) {
+func (self *BackoffDialPolicy) recordFailure() {
 	self.mu.Lock()
 	defer self.mu.Unlock()
-	s, ok := self.state[underlayType]
-	if !ok {
-		s = &dialState{}
-		self.state[underlayType] = s
-	}
-	s.consecutiveFailures++
-	s.lastSuccess = time.Time{}
+	self.consecutiveFailures++
+	self.lastSuccess = time.Time{}
 }
 
-// MarkStable resets the failure count for an underlay type.
-// Call this when a connection has been alive long enough to be considered stable.
-func (self *BackoffDialPolicy) MarkStable(underlayType string) {
+// MarkStable resets the failure count. Call this when a connection has been alive long
+// enough to be considered stable.
+func (self *BackoffDialPolicy) MarkStable() {
 	self.mu.Lock()
 	defer self.mu.Unlock()
-	if s, ok := self.state[underlayType]; ok {
-		s.consecutiveFailures = 0
-		s.lastSuccess = time.Time{}
-	}
+	self.consecutiveFailures = 0
+	self.lastSuccess = time.Time{}
 }
 
-// ConsecutiveFailures returns the number of consecutive failures for an underlay type.
-func (self *BackoffDialPolicy) ConsecutiveFailures(underlayType string) uint32 {
+// ConsecutiveFailures returns the current consecutive failure count.
+func (self *BackoffDialPolicy) ConsecutiveFailures() uint32 {
 	self.mu.Lock()
 	defer self.mu.Unlock()
-	if s, ok := self.state[underlayType]; ok {
-		return s.consecutiveFailures
-	}
-	return 0
+	return self.consecutiveFailures
 }
 
+// Dial attempts to create a new underlay, applying exponential backoff if previous attempts failed.
 func (self *BackoffDialPolicy) Dial(underlayType string, connectionId string, groupSecret []byte, connectTimeout time.Duration, cancel <-chan struct{}) (Underlay, error) {
 	log := pfxlog.Logger().WithField("underlayType", underlayType).WithField("connectionId", connectionId)
 
-	self.recordShortLivedConnection(underlayType)
+	self.recordShortLivedConnection()
 
-	if delay := self.getBackoffDelay(underlayType); delay > 0 {
+	if delay := self.getBackoffDelay(); delay > 0 {
 		log.WithField("delay", delay).Debug("backing off before dial")
 		select {
 		case <-cancel:
@@ -177,11 +149,11 @@ func (self *BackoffDialPolicy) Dial(underlayType string, connectionId string, gr
 		IsGroupedHeader:    {1},
 	})
 	if err != nil {
-		self.recordFailure(underlayType)
+		self.recordFailure()
 		log.WithError(err).Info("dial failed")
 		return nil, err
 	}
 
-	self.recordSuccess(underlayType)
+	self.recordSuccess()
 	return underlay, nil
 }