| 1 | package framework |
| 2 | |
| 3 | import ( |
| 4 | "time" |
| 5 | ) |
| 6 | |
| 7 | // ProtocolClient provides automatic instrumentation for protocol operations |
| 8 | type ProtocolClient struct { |
| 9 | name string |
| 10 | metrics *ProtocolMetrics |
| 11 | state *CollectorState |
| 12 | backoff *ExponentialBackoff |
| 13 | connectionIteration int64 // Iteration when this protocol connected |
| 14 | } |
| 15 | |
| 16 | // NewProtocolClient creates a new instrumented protocol client |
| 17 | func NewProtocolClient(name string, state *CollectorState) *ProtocolClient { |
| 18 | return &ProtocolClient{ |
| 19 | name: name, |
| 20 | metrics: state.RegisterProtocol(name), |
| 21 | state: state, |
| 22 | backoff: NewExponentialBackoff(), |
| 23 | } |
| 24 | } |
| 25 | |
| 26 | // GetBackoff returns the exponential backoff for this protocol |
| 27 | func (p *ProtocolClient) GetBackoff() *ExponentialBackoff { |
| 28 | p.backoff.protocol = p // Set reference for logging |
| 29 | return p.backoff |
| 30 | } |
| 31 | |
| 32 | // MarkConnected should be called when the protocol successfully connects |
| 33 | func (p *ProtocolClient) MarkConnected() { |
| 34 | p.connectionIteration = p.state.GetIteration() |
| 35 | p.backoff.Reset() |
| 36 | } |
| 37 | |
| 38 | // Debugf logs a debug message prefixed with the protocol name |
| 39 | func (p *ProtocolClient) Debugf(format string, args ...any) { |
| 40 | p.state.Debugf("[%s] "+format, append([]any{p.name}, args...)...) |
| 41 | } |
| 42 | |
| 43 | // Warningf logs a warning message prefixed with the protocol name |
| 44 | func (p *ProtocolClient) Warningf(format string, args ...any) { |
| 45 | p.state.Warningf("[%s] "+format, append([]any{p.name}, args...)...) |
| 46 | } |
| 47 | |
| 48 | // Errorf logs an error message prefixed with the protocol name |
| 49 | func (p *ProtocolClient) Errorf(format string, args ...any) { |
| 50 | p.state.Errorf("[%s] "+format, append([]any{p.name}, args...)...) |
| 51 | } |
| 52 | |
| 53 | // Infof logs an info message prefixed with the protocol name |
| 54 | func (p *ProtocolClient) Infof(format string, args ...any) { |
| 55 | p.state.Infof("[%s] "+format, append([]any{p.name}, args...)...) |
| 56 | } |
| 57 | |
| 58 | // IsReconnect returns true if this is the same iteration as when connected |
| 59 | func (p *ProtocolClient) IsReconnect() bool { |
| 60 | return p.connectionIteration == p.state.GetIteration() |
| 61 | } |
| 62 | |
| 63 | // Track wraps a protocol operation with automatic metrics collection |
| 64 | func (p *ProtocolClient) Track(operationName string, fn func() error) error { |
| 65 | start := time.Now() |
| 66 | |
| 67 | // Execute the operation |
| 68 | err := fn() |
| 69 | |
| 70 | // Track metrics |
| 71 | elapsed := time.Since(start).Microseconds() |
| 72 | p.metrics.RequestCount++ |
| 73 | p.metrics.TotalLatency += elapsed |
| 74 | |
| 75 | if elapsed > p.metrics.MaxLatency { |
| 76 | p.metrics.MaxLatency = elapsed |
| 77 | } |
| 78 | |
| 79 | if err != nil { |
| 80 | p.metrics.ErrorCount++ |
| 81 | } |
| 82 | |
| 83 | return err |
| 84 | } |
| 85 | |
| 86 | // TrackWithSize wraps operations that have request/response sizes |
| 87 | func (p *ProtocolClient) TrackWithSize(operationName string, requestSize int64, fn func() (int64, error)) error { |
| 88 | start := time.Now() |
| 89 | |
| 90 | // Track request size |
| 91 | p.metrics.BytesSent += requestSize |
| 92 | |
| 93 | // Execute the operation |
| 94 | responseSize, err := fn() |
| 95 | |
| 96 | // Track response size |
| 97 | if responseSize > 0 { |
| 98 | p.metrics.BytesReceived += responseSize |
| 99 | } |
| 100 | |
| 101 | // Track timing |
| 102 | elapsed := time.Since(start).Microseconds() |
| 103 | p.metrics.RequestCount++ |
| 104 | p.metrics.TotalLatency += elapsed |
| 105 | |
| 106 | if elapsed > p.metrics.MaxLatency { |
| 107 | p.metrics.MaxLatency = elapsed |
| 108 | } |
| 109 | |
| 110 | if err != nil { |
| 111 | p.metrics.ErrorCount++ |
| 112 | } |
| 113 | |
| 114 | return err |
| 115 | } |
| 116 | |
| 117 | // GetProtocolMetrics returns current protocol metrics for charting |
| 118 | func (p *ProtocolClient) GetProtocolMetrics() map[string]int64 { |
| 119 | metrics := make(map[string]int64) |
| 120 | |
| 121 | // Operation counts |
| 122 | metrics["requests"] = p.metrics.RequestCount |
| 123 | metrics["errors"] = p.metrics.ErrorCount |
| 124 | |
| 125 | // Latency metrics |
| 126 | if p.metrics.RequestCount > 0 { |
| 127 | metrics["avg_latency"] = p.metrics.TotalLatency / p.metrics.RequestCount |
| 128 | } |
| 129 | metrics["max_latency"] = p.metrics.MaxLatency |
| 130 | |
| 131 | // Size metrics |
| 132 | metrics["bytes_sent"] = p.metrics.BytesSent |
| 133 | metrics["bytes_received"] = p.metrics.BytesReceived |
| 134 | |
| 135 | return metrics |
| 136 | } |
| 137 | |
| 138 | // ClassifyError wraps errors with retry classification |
| 139 | func ClassifyError(err error, errType ErrorType) error { |
| 140 | if err == nil { |
| 141 | return nil |
| 142 | } |
| 143 | return CollectorError{Err: err, Type: errType} |
| 144 | } |
| 145 | |
| 146 | // IsTemporary checks if an error is temporary and should be retried |
| 147 | func IsTemporary(err error) bool { |
| 148 | if collErr, ok := err.(CollectorError); ok { |
| 149 | return collErr.Type == ErrorTemporary |
| 150 | } |
| 151 | return false |
| 152 | } |
| 153 | |
| 154 | // IsFatal checks if an error is fatal and module should be disabled |
| 155 | func IsFatal(err error) bool { |
| 156 | if collErr, ok := err.(CollectorError); ok { |
| 157 | return collErr.Type == ErrorFatal || collErr.Type == ErrorAuth |
| 158 | } |
| 159 | return false |
| 160 | } |
| 161 | |
| 162 | // ExponentialBackoff implements retry logic with backoff |
| 163 | type ExponentialBackoff struct { |
| 164 | InitialInterval time.Duration |
| 165 | MaxInterval time.Duration |
| 166 | Multiplier float64 |
| 167 | currentInterval time.Duration |
| 168 | attempt int |
| 169 | protocol *ProtocolClient // For logging |
| 170 | } |
| 171 | |
| 172 | // NewExponentialBackoff creates a new backoff handler |
| 173 | func NewExponentialBackoff() *ExponentialBackoff { |
| 174 | return &ExponentialBackoff{ |
| 175 | InitialInterval: time.Second, |
| 176 | MaxInterval: 5 * time.Minute, |
| 177 | Multiplier: 2.0, |
| 178 | currentInterval: time.Second, |
| 179 | } |
| 180 | } |
| 181 | |
| 182 | // NextInterval returns the next backoff interval |
| 183 | func (b *ExponentialBackoff) NextInterval() time.Duration { |
| 184 | defer func() { |
| 185 | b.currentInterval = min(time.Duration(float64(b.currentInterval)*b.Multiplier), b.MaxInterval) |
| 186 | b.attempt++ |
| 187 | }() |
| 188 | |
| 189 | // Add jitter (±10%) |
| 190 | jitter := time.Duration(float64(b.currentInterval) * 0.1) |
| 191 | return b.currentInterval + jitter |
| 192 | } |
| 193 | |
| 194 | // Reset resets the backoff to initial state |
| 195 | func (b *ExponentialBackoff) Reset() { |
| 196 | b.currentInterval = b.InitialInterval |
| 197 | b.attempt = 0 |
| 198 | } |
| 199 | |
| 200 | // ShouldRetry determines if we should retry based on attempt count |
| 201 | func (b *ExponentialBackoff) ShouldRetry() bool { |
| 202 | // Retry up to 10 times (reaches max interval after ~8 attempts) |
| 203 | return b.attempt < 2 |
| 204 | } |
| 205 | |
| 206 | // Retry executes a function with exponential backoff |
| 207 | func (b *ExponentialBackoff) Retry(fn func() error) error { |
| 208 | var lastErr error |
| 209 | |
| 210 | if b.protocol != nil { |
| 211 | b.protocol.Debugf("Starting retry sequence with initial interval %v, max interval %v", |
| 212 | b.InitialInterval, b.MaxInterval) |
| 213 | } |
| 214 | |
| 215 | for { |
| 216 | if b.protocol != nil && b.attempt > 0 { |
| 217 | b.protocol.Debugf("Retry attempt #%d (after %d failed attempts)", b.attempt+1, b.attempt) |
| 218 | } |
| 219 | |
| 220 | err := fn() |
| 221 | if err == nil { |
| 222 | if b.protocol != nil && b.attempt > 0 { |
| 223 | b.protocol.Infof("Operation succeeded after %d retry attempts", b.attempt) |
| 224 | } |
| 225 | b.Reset() |
| 226 | return nil |
| 227 | } |
| 228 | |
| 229 | lastErr = err |
| 230 | |
| 231 | // If this is a fatal error, don't retry |
| 232 | if IsFatal(err) { |
| 233 | if b.protocol != nil { |
| 234 | b.protocol.Errorf("Fatal error encountered, stopping retries: %v", err) |
| 235 | } |
| 236 | return err |
| 237 | } |
| 238 | |
| 239 | // Check if we should continue retrying |
| 240 | if !b.ShouldRetry() { |
| 241 | if b.protocol != nil { |
| 242 | b.protocol.Errorf("Maximum retry attempts (%d) reached, giving up: %v", b.attempt, lastErr) |
| 243 | } |
| 244 | return lastErr |
| 245 | } |
| 246 | |
| 247 | // Wait before next attempt |
| 248 | interval := b.NextInterval() |
| 249 | if b.protocol != nil { |
| 250 | b.protocol.Debugf("Waiting %v before retry attempt #%d (error was: %v)", |
| 251 | interval, b.attempt+1, err) |
| 252 | } |
| 253 | time.Sleep(interval) |
| 254 | } |
| 255 | } |