| 1 | package main |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "crypto/sha256" |
| 6 | "encoding/json" |
| 7 | "errors" |
| 8 | "expvar" |
| 9 | "flag" |
| 10 | "fmt" |
| 11 | "io" |
| 12 | "log/slog" |
| 13 | "net/http" |
| 14 | _ "net/http/pprof" // Import for side effects - registers HTTP handlers |
| 15 | "os" |
| 16 | "os/signal" |
| 17 | "runtime" |
| 18 | "strings" |
| 19 | "sync" |
| 20 | "sync/atomic" |
| 21 | "syscall" |
| 22 | "time" |
| 23 | |
| 24 | // Import promhttp for the handler |
| 25 | "github.com/prometheus/client_golang/prometheus/promhttp" |
| 26 | "github.com/tidwall/gjson" |
| 27 | "go.opentelemetry.io/otel" |
| 28 | "go.opentelemetry.io/otel/attribute" |
| 29 | "go.opentelemetry.io/otel/exporters/prometheus" |
| 30 | "go.opentelemetry.io/otel/metric" |
| 31 | sdkmetric "go.opentelemetry.io/otel/sdk/metric" |
| 32 | ) |
| 33 | |
| 34 | // --- Constants --- |
| 35 | const ( |
| 36 | maxRequestBodySize = 20 * 1024 // 20 KiB |
| 37 | ) |
| 38 | |
| 39 | // --- Custom Flag Type for Multi-use --dedup-key --- |
| 40 | type dedupPaths []string |
| 41 | |
| 42 | func (d *dedupPaths) String() string { return fmt.Sprintf("%v", *d) } |
| 43 | func (d *dedupPaths) Set(value string) error { |
| 44 | if value == "" { |
| 45 | return fmt.Errorf("dedup-key path cannot be empty") |
| 46 | } |
| 47 | *d = append(*d, value) |
| 48 | return nil |
| 49 | } |
| 50 | |
| 51 | // --- Global variables --- |
| 52 | var ( |
| 53 | // Core functionality variables |
| 54 | seenIDs map[[32]byte]seenEntry |
| 55 | mapMutex = &sync.Mutex{} |
| 56 | dedupWindow time.Duration |
| 57 | keyPaths dedupPaths |
| 58 | dedupSeparator string |
| 59 | startTime time.Time |
| 60 | dedupLogger *os.File |
| 61 | dedupLogFile string // Store the deduplication log file path |
| 62 | |
| 63 | // Track active connections for graceful shutdown |
| 64 | activeConnections int32 |
| 65 | |
| 66 | // OpenTelemetry meter provider and metrics |
| 67 | meterProvider *sdkmetric.MeterProvider |
| 68 | meter metric.Meter |
| 69 | |
| 70 | // Counters |
| 71 | requestsCounter metric.Int64Counter // Unified counter with status label |
| 72 | bytesReceived metric.Int64Counter |
| 73 | |
| 74 | // Gauges |
| 75 | dedupCacheSize metric.Int64ObservableGauge |
| 76 | activeConnectionsGauge metric.Int64ObservableGauge |
| 77 | uptimeGauge metric.Int64ObservableGauge |
| 78 | |
| 79 | // Histograms |
| 80 | requestDuration metric.Float64Histogram |
| 81 | |
| 82 | // Keep expvar for backward compatibility |
| 83 | eventMetrics = expvar.NewMap("agent_events") |
| 84 | ) |
| 85 | |
| 86 | // --- Data Structures --- |
| 87 | type seenEntry struct { |
| 88 | timestamp time.Time |
| 89 | } |
| 90 | |
| 91 | // --- Core Logic Functions --- |
| 92 | |
| 93 | // initMetrics initializes OpenTelemetry metrics and expvar metrics for backward compatibility |
| 94 | func initMetrics() (*prometheus.Exporter, error) { |
| 95 | // Set up the OpenTelemetry Prometheus exporter |
| 96 | exporter, err := prometheus.New() |
| 97 | if err != nil { |
| 98 | // Return error instead of exiting directly |
| 99 | return nil, fmt.Errorf("failed to create Prometheus exporter: %w", err) |
| 100 | } |
| 101 | |
| 102 | // Create a new meter provider with the Prometheus exporter as a Reader |
| 103 | meterProvider = sdkmetric.NewMeterProvider( |
| 104 | sdkmetric.WithReader(exporter), |
| 105 | ) |
| 106 | otel.SetMeterProvider(meterProvider) |
| 107 | |
| 108 | // Create a new meter |
| 109 | meter = meterProvider.Meter("agent-events") |
| 110 | |
| 111 | // --- Create Metric Instruments --- |
| 112 | // Helper to create counters with status labels for consolidated metrics |
| 113 | createLabeledCounter := func(name, desc string) (metric.Int64Counter, error) { |
| 114 | counter, err := meter.Int64Counter( |
| 115 | name, |
| 116 | metric.WithDescription(desc), |
| 117 | metric.WithUnit("1"), |
| 118 | ) |
| 119 | if err != nil { |
| 120 | slog.Error("failed to create labeled counter", "name", name, "error", err) |
| 121 | } |
| 122 | return counter, err |
| 123 | } |
| 124 | createGauge := func(name, desc string) (metric.Int64ObservableGauge, error) { |
| 125 | gauge, err := meter.Int64ObservableGauge(name, metric.WithDescription(desc)) |
| 126 | if err != nil { |
| 127 | slog.Error("failed to create gauge", "name", name, "error", err) |
| 128 | } |
| 129 | return gauge, err |
| 130 | } |
| 131 | createHistogram := func(name, desc string, buckets []float64) (metric.Float64Histogram, error) { |
| 132 | hist, err := meter.Float64Histogram(name, |
| 133 | metric.WithDescription(desc), |
| 134 | metric.WithExplicitBucketBoundaries(buckets...), |
| 135 | ) |
| 136 | if err != nil { |
| 137 | slog.Error("failed to create histogram", "name", name, "error", err) |
| 138 | } |
| 139 | return hist, err |
| 140 | } |
| 141 | |
| 142 | // Create unified counters with status label |
| 143 | requestsCounter, _ = createLabeledCounter("agent_events_requests", "Number of requests by status") |
| 144 | bytesReceived, _ = createLabeledCounter("agent_events_received_bytes", "Number of bytes received in request bodies by status") |
| 145 | |
| 146 | // Pre-initialize counters with all status labels set to zero |
| 147 | ctx := context.Background() |
| 148 | statusLabels := []string{"success", "duplicate", "invalid_json", "method_not_allowed", "body_too_large", "failed_to_read", "cant_marshal_output"} |
| 149 | for _, status := range statusLabels { |
| 150 | requestsCounter.Add(ctx, 0, metric.WithAttributes(attribute.String("status", status))) |
| 151 | bytesReceived.Add(ctx, 0, metric.WithAttributes(attribute.String("status", status))) |
| 152 | } |
| 153 | |
| 154 | dedupCacheSize, _ = createGauge("agent_events_dedup_cache_entries", "Current number of entries in the deduplication cache") |
| 155 | activeConnectionsGauge, _ = createGauge("agent_events_active_connections", "Number of currently active connections") |
| 156 | uptimeGauge, _ = createGauge("agent_events_uptime_seconds", "How long the server has been running in seconds") |
| 157 | requestDuration, _ = createHistogram("agent_events_request_duration_seconds", |
| 158 | "Histogram of request processing times in seconds", |
| 159 | []float64{0.00001, 0.000025, 0.00005, 0.0001, 0.00025, 0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1}, |
| 160 | ) |
| 161 | // Basic check if any metric failed (optional, depends on how critical individual metrics are) |
| 162 | // if requestsTotal == nil || ... { return exporter, fmt.Errorf("one or more metrics failed to initialize") } |
| 163 | |
| 164 | // Register callbacks for observable metrics |
| 165 | _, err = meter.RegisterCallback( |
| 166 | func(_ context.Context, observer metric.Observer) error { |
| 167 | mapMutex.Lock() |
| 168 | observer.ObserveInt64(dedupCacheSize, int64(len(seenIDs))) |
| 169 | mapMutex.Unlock() |
| 170 | observer.ObserveInt64(activeConnectionsGauge, int64(atomic.LoadInt32(&activeConnections))) |
| 171 | observer.ObserveInt64(uptimeGauge, int64(time.Since(startTime).Seconds())) |
| 172 | return nil |
| 173 | }, |
| 174 | dedupCacheSize, activeConnectionsGauge, uptimeGauge, |
| 175 | ) |
| 176 | if err != nil { |
| 177 | // Log but don't necessarily exit, maybe gauges won't update |
| 178 | slog.Error("failed to register callback for observable metrics", "error", err) |
| 179 | } |
| 180 | |
| 181 | // Create expvar metrics for backward compatibility |
| 182 | eventMetrics.Set("active_connections", expvar.Func(func() interface{} { return atomic.LoadInt32(&activeConnections) })) |
| 183 | eventMetrics.Set("uptime_seconds", expvar.Func(func() interface{} { return time.Since(startTime).Seconds() })) |
| 184 | eventMetrics.Set("start_time_seconds", expvar.Func(func() interface{} { return startTime.Unix() })) |
| 185 | eventMetrics.Set("runtime_stats", expvar.Func(func() interface{} { |
| 186 | memStats := &runtime.MemStats{} |
| 187 | runtime.ReadMemStats(memStats) |
| 188 | return *memStats |
| 189 | })) |
| 190 | |
| 191 | return exporter, nil // Return the exporter (even though we use promhttp.Handler) and nil error |
| 192 | } |
| 193 | |
| 194 | // checkAndRecordHash accepts the SHA256 hash ([32]byte) for checking. |
| 195 | // It now REFRESHES the timestamp whenever a hash is found, |
| 196 | // effectively creating a sliding deduplication window. |
| 197 | func checkAndRecordHash(hash [32]byte) bool { |
| 198 | now := time.Now() |
| 199 | mapMutex.Lock() |
| 200 | defer mapMutex.Unlock() |
| 201 | var zeroHash [32]byte |
| 202 | if hash == zeroHash { |
| 203 | slog.Warn("potentially zero hash received", "hash", fmt.Sprintf("%x", hash)) |
| 204 | } |
| 205 | entry, found := seenIDs[hash] |
| 206 | if found { |
| 207 | isRecentDuplicate := now.Sub(entry.timestamp) < dedupWindow |
| 208 | seenIDs[hash] = seenEntry{timestamp: now} |
| 209 | return !isRecentDuplicate |
| 210 | } |
| 211 | seenIDs[hash] = seenEntry{timestamp: now} |
| 212 | return true |
| 213 | } |
| 214 | |
| 215 | // cleanupExpiredEntries uses the hash ([32]byte) as the key type. |
| 216 | func cleanupExpiredEntries(interval time.Duration) { |
| 217 | ticker := time.NewTicker(interval) |
| 218 | defer ticker.Stop() |
| 219 | cleanedCount := 0 |
| 220 | lastCleanupLogTime := time.Now() |
| 221 | for range ticker.C { |
| 222 | mapMutex.Lock() |
| 223 | now := time.Now() |
| 224 | currentMapSize := len(seenIDs) |
| 225 | deletedInCycle := 0 // Track deletes per cycle for more granular debug |
| 226 | for h, entry := range seenIDs { |
| 227 | if now.Sub(entry.timestamp) >= dedupWindow { |
| 228 | delete(seenIDs, h) |
| 229 | cleanedCount++ |
| 230 | deletedInCycle++ |
| 231 | } |
| 232 | } |
| 233 | mapMutex.Unlock() |
| 234 | |
| 235 | // Log hourly summary if any were cleaned in the last hour |
| 236 | if cleanedCount > 0 && time.Since(lastCleanupLogTime) >= time.Hour { |
| 237 | slog.Debug("cleaned up expired entries", |
| 238 | "count_past_hour", cleanedCount, |
| 239 | "remaining_entries", currentMapSize-deletedInCycle) // Use size before delete for consistency |
| 240 | cleanedCount = 0 // Reset hourly count |
| 241 | lastCleanupLogTime = time.Now() |
| 242 | } else if deletedInCycle > 0 { |
| 243 | // Optional: Log every cycle if debugging cleanup |
| 244 | // slog.Debug("cleanup cycle completed", "deleted", deletedInCycle, "remaining", currentMapSize-deletedInCycle) |
| 245 | } |
| 246 | } |
| 247 | } |
| 248 | |
| 249 | // connectionTracker is middleware that wraps an http.Handler to track active connections |
| 250 | func connectionTracker(next http.Handler) http.Handler { |
| 251 | return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 252 | atomic.AddInt32(&activeConnections, 1) |
| 253 | defer atomic.AddInt32(&activeConnections, -1) |
| 254 | next.ServeHTTP(w, r) |
| 255 | }) |
| 256 | } |
| 257 | |
| 258 | // --- HTTP Handler --- |
| 259 | func handler(w http.ResponseWriter, r *http.Request) { |
| 260 | requestStartTime := time.Now() |
| 261 | ctx := r.Context() |
| 262 | |
| 263 | defer func() { |
| 264 | requestDuration.Record(ctx, time.Since(requestStartTime).Seconds()) |
| 265 | }() |
| 266 | |
| 267 | // Method Check |
| 268 | if r.Method != http.MethodPost { |
| 269 | http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed) |
| 270 | slog.Info("request discarded", "reason", "method_not_allowed", "method", r.Method, "remote_addr", r.RemoteAddr) |
| 271 | requestsCounter.Add(ctx, 1, metric.WithAttributes(attribute.String("status", "method_not_allowed"))) |
| 272 | return |
| 273 | } |
| 274 | |
| 275 | // Read Body & Size Check |
| 276 | r.Body = http.MaxBytesReader(w, r.Body, maxRequestBodySize) |
| 277 | body, err := io.ReadAll(r.Body) |
| 278 | if err != nil { |
| 279 | var maxBytesErr *http.MaxBytesError |
| 280 | if errors.As(err, &maxBytesErr) { |
| 281 | http.Error(w, fmt.Sprintf("Request body exceeds limit (%d bytes)", maxRequestBodySize), http.StatusRequestEntityTooLarge) |
| 282 | slog.Info("request discarded", "reason", "body_too_large", "limit", maxRequestBodySize, "remote_addr", r.RemoteAddr) |
| 283 | requestsCounter.Add(ctx, 1, metric.WithAttributes(attribute.String("status", "body_too_large"))) |
| 284 | } else { |
| 285 | http.Error(w, "Error reading request", http.StatusInternalServerError) |
| 286 | slog.Error("request discarded", "reason", "error_reading_body", "error", err) |
| 287 | requestsCounter.Add(ctx, 1, metric.WithAttributes(attribute.String("status", "failed_to_read"))) |
| 288 | } |
| 289 | return |
| 290 | } |
| 291 | |
| 292 | // JSON Validation - Attempt to unmarshal directly to map (requires object) |
| 293 | var fullData map[string]interface{} |
| 294 | if err := json.Unmarshal(body, &fullData); err != nil { |
| 295 | http.Error(w, "Invalid JSON", http.StatusBadRequest) |
| 296 | bodyDetail := "" |
| 297 | if slog.Default().Enabled(context.Background(), slog.LevelDebug) { |
| 298 | bodyDetail = fmt.Sprintf(", Body: %s", string(body)) |
| 299 | } else { |
| 300 | bodyDetail = fmt.Sprintf(", Body snippet: %s", limitString(string(body), 100)) |
| 301 | } |
| 302 | slog.Warn("request discarded", "reason", "invalid_json", "error", err.Error(), "body_detail", bodyDetail) |
| 303 | bytesReceived.Add(ctx, int64(len(body)), metric.WithAttributes(attribute.String("status", "invalid_json"))) |
| 304 | requestsCounter.Add(ctx, 1, metric.WithAttributes(attribute.String("status", "invalid_json"))) |
| 305 | return |
| 306 | } |
| 307 | |
| 308 | bytesReceived.Add(ctx, int64(len(body)), metric.WithAttributes(attribute.String("status", "success"))) |
| 309 | |
| 310 | // Add Cloudflare Headers to all requests, excluding IP addresses for GDPR compliance |
| 311 | cfHeaders := make(map[string]string) |
| 312 | cfHeaderPrefixes := []string{"CF-IPCountry", "CF-Ray", "CF-IPCity", "CF-IPContinent", "CF-IPRegion", "CF-IPTimeZone", "CF-IPCOLO"} |
| 313 | // Explicitly excluding IP-related headers: CF-Connecting-IP, CF-IPLatitude, CF-IPLongitude, CF-Visitor |
| 314 | for _, name := range cfHeaderPrefixes { |
| 315 | if value := r.Header.Get(name); value != "" { |
| 316 | key := strings.TrimPrefix(name, "CF-") |
| 317 | cfHeaders[key] = value |
| 318 | } |
| 319 | } |
| 320 | for name, values := range r.Header { |
| 321 | if strings.HasPrefix(name, "CF-") && len(values) > 0 { |
| 322 | // Skip IP-related headers for GDPR compliance |
| 323 | if name == "CF-Connecting-IP" || name == "CF-IPLatitude" || name == "CF-IPLongitude" || name == "CF-Visitor" { |
| 324 | continue |
| 325 | } |
| 326 | key := strings.TrimPrefix(name, "CF-") |
| 327 | if _, exists := cfHeaders[key]; !exists { |
| 328 | cfHeaders[key] = values[0] |
| 329 | } |
| 330 | } |
| 331 | } |
| 332 | if len(cfHeaders) > 0 { |
| 333 | fullData["cf"] = cfHeaders |
| 334 | slog.Debug("added cloudflare headers", "count", len(cfHeaders)) |
| 335 | } |
| 336 | |
| 337 | // Deduplication Logic |
| 338 | shouldProcess := true |
| 339 | var finalKeyString string |
| 340 | var dedupHash [32]byte |
| 341 | if len(keyPaths) > 0 { |
| 342 | var keyBuilder strings.Builder |
| 343 | for i, path := range keyPaths { |
| 344 | result := gjson.GetBytes(body, path) |
| 345 | keyBuilder.WriteString(result.String()) // gjson returns "" for non-existent paths |
| 346 | if i < len(keyPaths)-1 { |
| 347 | keyBuilder.WriteString(dedupSeparator) |
| 348 | } |
| 349 | } |
| 350 | finalKeyString = keyBuilder.String() |
| 351 | dedupHash = sha256.Sum256([]byte(finalKeyString)) |
| 352 | slog.Debug("generated dedup key", "key_string", finalKeyString, "hash", fmt.Sprintf("%x", dedupHash)) |
| 353 | |
| 354 | // Add _dedup key to all requests (regardless of duplicate status) |
| 355 | fullData["_dedup"] = map[string]interface{}{ |
| 356 | "key": finalKeyString, |
| 357 | "hash": fmt.Sprintf("%x", dedupHash), |
| 358 | } |
| 359 | |
| 360 | // Check if this is a duplicate |
| 361 | if !checkAndRecordHash(dedupHash) { |
| 362 | shouldProcess = false |
| 363 | bytesReceived.Add(ctx, int64(len(body)), metric.WithAttributes(attribute.String("status", "duplicate"))) |
| 364 | requestsCounter.Add(ctx, 1, metric.WithAttributes(attribute.String("status", "duplicate"))) |
| 365 | slog.Debug("discarded duplicate request", "hash", fmt.Sprintf("%x", dedupHash), "note", "timestamp refreshed") |
| 366 | } |
| 367 | } else { |
| 368 | slog.Debug("skipping deduplication", "reason", "no dedup keys provided") |
| 369 | } |
| 370 | |
| 371 | // Marshal the fully prepared object for either stdout or dedup log |
| 372 | outputBytes, err := json.Marshal(fullData) |
| 373 | if err != nil { |
| 374 | http.Error(w, "Internal Server Error during output marshal", http.StatusInternalServerError) |
| 375 | slog.Error("request discarded", "reason", "json_marshal_failed", "error", err) |
| 376 | requestsCounter.Add(ctx, 1, metric.WithAttributes(attribute.String("status", "cant_marshal_output"))) |
| 377 | return |
| 378 | } |
| 379 | |
| 380 | // Decide where to output based on deduplication status |
| 381 | if shouldProcess { |
| 382 | // Write to stdout for normal processing |
| 383 | fmt.Println(string(outputBytes)) |
| 384 | requestsCounter.Add(ctx, 1, metric.WithAttributes(attribute.String("status", "success"))) |
| 385 | } else if dedupLogger != nil { |
| 386 | // Write to dedup log file if it's a duplicate and logging is enabled |
| 387 | if _, err := fmt.Fprintln(dedupLogger, string(outputBytes)); err != nil { |
| 388 | slog.Error("failed to write to deduplication log file", "error", err) |
| 389 | } |
| 390 | } |
| 391 | |
| 392 | // Send response |
| 393 | if _, err := w.Write([]byte("OK")); err != nil { |
| 394 | if shouldProcess { |
| 395 | slog.Error("error writing response", "context", "after_successful_processing", "error", err) |
| 396 | } else { |
| 397 | slog.Error("error writing response", "context", "after_duplicate_discard", "error", err) |
| 398 | } |
| 399 | } |
| 400 | |
| 401 | // For duplicates, return early |
| 402 | if !shouldProcess { |
| 403 | return |
| 404 | } |
| 405 | } |
| 406 | |
| 407 | // healthHandler provides a simple health check endpoint |
| 408 | func healthHandler(w http.ResponseWriter, r *http.Request) { |
| 409 | status := map[string]interface{}{ |
| 410 | "status": "ok", |
| 411 | "timestamp": time.Now().Format(time.RFC3339), |
| 412 | "uptime": time.Since(startTime).String(), |
| 413 | "goroutines": runtime.NumGoroutine(), |
| 414 | "connections": atomic.LoadInt32(&activeConnections), |
| 415 | "version": "agent-events v1.0", |
| 416 | } |
| 417 | memStats := &runtime.MemStats{} |
| 418 | runtime.ReadMemStats(memStats) |
| 419 | status["memory"] = map[string]interface{}{"alloc": memStats.Alloc, "total_alloc": memStats.TotalAlloc, "sys": memStats.Sys, "heap_alloc": memStats.HeapAlloc, "gc_cycles": memStats.NumGC} |
| 420 | mapMutex.Lock() |
| 421 | mapSize := len(seenIDs) |
| 422 | mapMutex.Unlock() |
| 423 | dedupLogEnabled := dedupLogger != nil |
| 424 | dedupLogPath := "" |
| 425 | if dedupLogEnabled { |
| 426 | dedupLogPath = dedupLogFile |
| 427 | } |
| 428 | status["deduplication"] = map[string]interface{}{ |
| 429 | "enabled": len(keyPaths) > 0, |
| 430 | "keys": keyPaths, |
| 431 | "window": dedupWindow.String(), |
| 432 | "map_size": mapSize, |
| 433 | "log_enabled": dedupLogEnabled, |
| 434 | "log_file": dedupLogPath, |
| 435 | } |
| 436 | w.Header().Set("Content-Type", "application/json") |
| 437 | w.WriteHeader(http.StatusOK) |
| 438 | if err := json.NewEncoder(w).Encode(status); err != nil { |
| 439 | slog.Error("error encoding health check response", "error", err) |
| 440 | } |
| 441 | } |
| 442 | |
| 443 | // main is the entry point of the application |
| 444 | func main() { |
| 445 | // Setup initial logger before flag parsing |
| 446 | initialLogLevel := slog.LevelInfo |
| 447 | initialLogHandler := slog.NewJSONHandler(os.Stderr, &slog.HandlerOptions{Level: initialLogLevel, AddSource: true}) |
| 448 | slog.SetDefault(slog.New(initialLogHandler)) |
| 449 | |
| 450 | startTime = time.Now() |
| 451 | |
| 452 | // Define flags |
| 453 | port := flag.Int("port", 8080, "Port to listen on") |
| 454 | dedupSeconds := flag.Int("dedup-window", 1800, "Deduplication window in seconds") |
| 455 | metricsPath := flag.String("metrics-path", "/metrics", "Path for OpenTelemetry Prometheus metrics endpoint") |
| 456 | expvarPath := flag.String("expvar-path", "/debug/vars", "Path for expvar metrics endpoint (empty to disable)") |
| 457 | healthPath := flag.String("health-path", "/healthz", "Path for health check endpoint") |
| 458 | logFormat := flag.String("log-format", "json", "Log format: 'json' or 'text'") |
| 459 | logLevelFlag := flag.String("log-level", "info", "Log level: 'debug', 'info', 'warn', 'error'") |
| 460 | // Use the global dedupLogFile variable |
| 461 | flag.StringVar(&dedupLogFile, "dedup-logfile", "", "File to log deduplicated requests (empty to disable)") |
| 462 | flag.Var(&keyPaths, "dedup-key", "JSON path (dot-notation) for deduplication key (can be used multiple times)") |
| 463 | flag.StringVar(&dedupSeparator, "dedup-separator", "-", "Separator used between multi-key values") |
| 464 | flag.Parse() |
| 465 | |
| 466 | // Configure final logger based on flags |
| 467 | var level slog.Level |
| 468 | switch strings.ToLower(*logLevelFlag) { |
| 469 | case "debug": |
| 470 | level = slog.LevelDebug |
| 471 | case "info": |
| 472 | level = slog.LevelInfo |
| 473 | case "warn": |
| 474 | level = slog.LevelWarn |
| 475 | case "error": |
| 476 | level = slog.LevelError |
| 477 | default: |
| 478 | slog.Warn("invalid log level specified, defaulting to info", "value", *logLevelFlag) |
| 479 | level = slog.LevelInfo |
| 480 | } |
| 481 | var logHandler slog.Handler |
| 482 | if *logFormat == "text" { |
| 483 | logHandler = slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: level, AddSource: true}) |
| 484 | } else { |
| 485 | logHandler = slog.NewJSONHandler(os.Stderr, &slog.HandlerOptions{Level: level, AddSource: true}) |
| 486 | } |
| 487 | slog.SetDefault(slog.New(logHandler)) |
| 488 | |
| 489 | // Initialize core components |
| 490 | seenIDs = make(map[[32]byte]seenEntry) |
| 491 | dedupWindow = time.Duration(*dedupSeconds) * time.Second |
| 492 | |
| 493 | // Open deduplication log file if specified |
| 494 | if dedupLogFile != "" { |
| 495 | var err error |
| 496 | dedupLogger, err = os.OpenFile(dedupLogFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644) |
| 497 | if err != nil { |
| 498 | slog.Error("failed to open deduplication log file", "file", dedupLogFile, "error", err) |
| 499 | os.Exit(1) |
| 500 | } |
| 501 | slog.Info("deduplication log file opened", "path", dedupLogFile) |
| 502 | } |
| 503 | |
| 504 | if _, err := initMetrics(); err != nil { // Handle potential error from initMetrics |
| 505 | slog.Error("failed to initialize metrics", "error", err) |
| 506 | os.Exit(1) |
| 507 | } |
| 508 | |
| 509 | // Start background tasks |
| 510 | if dedupWindow > 0 && len(keyPaths) > 0 { |
| 511 | cleanupInterval := dedupWindow / 10 |
| 512 | if cleanupInterval < 1*time.Minute { |
| 513 | cleanupInterval = 1 * time.Minute |
| 514 | } else if cleanupInterval > 15*time.Minute { |
| 515 | cleanupInterval = 15 * time.Minute |
| 516 | } |
| 517 | slog.Info("cleanup goroutine started", "interval", cleanupInterval) |
| 518 | go cleanupExpiredEntries(cleanupInterval) |
| 519 | } else if dedupWindow <= 0 && len(keyPaths) > 0 { |
| 520 | slog.Warn("deduplication keys provided, but window is zero or negative", "keys", keyPaths, "window", dedupWindow) |
| 521 | } |
| 522 | |
| 523 | // Configure HTTP server |
| 524 | server := &http.Server{ |
| 525 | Addr: fmt.Sprintf(":%d", *port), |
| 526 | ReadTimeout: 10 * time.Second, |
| 527 | WriteTimeout: 10 * time.Second, |
| 528 | IdleTimeout: 60 * time.Second, |
| 529 | } |
| 530 | mux := http.NewServeMux() |
| 531 | mux.HandleFunc("/", handler) |
| 532 | mux.HandleFunc(*healthPath, healthHandler) |
| 533 | if *expvarPath != "" { |
| 534 | mux.Handle(*expvarPath, expvar.Handler()) |
| 535 | } // Register expvar if path not empty |
| 536 | mux.Handle(*metricsPath, promhttp.Handler()) // Use promhttp handler for OTEL metrics |
| 537 | // Add pprof handlers to custom mux |
| 538 | mux.HandleFunc("/debug/pprof/", http.DefaultServeMux.ServeHTTP) |
| 539 | mux.HandleFunc("/debug/pprof/cmdline", http.DefaultServeMux.ServeHTTP) |
| 540 | mux.HandleFunc("/debug/pprof/profile", http.DefaultServeMux.ServeHTTP) |
| 541 | mux.HandleFunc("/debug/pprof/symbol", http.DefaultServeMux.ServeHTTP) |
| 542 | mux.HandleFunc("/debug/pprof/trace", http.DefaultServeMux.ServeHTTP) |
| 543 | server.Handler = connectionTracker(mux) // Apply middleware |
| 544 | |
| 545 | // Start server and handle shutdown |
| 546 | stop := make(chan os.Signal, 1) |
| 547 | signal.Notify(stop, os.Interrupt, syscall.SIGTERM) |
| 548 | serverErrors := make(chan error, 1) |
| 549 | |
| 550 | slog.Info("server starting", "port", *port, "metrics_path", *metricsPath, "expvar_path", *expvarPath, "health_path", *healthPath, "log_level", level.String()) // Simplified startup log |
| 551 | |
| 552 | go func() { |
| 553 | slog.Info("server listening", "addr", server.Addr) |
| 554 | serverErrors <- server.ListenAndServe() |
| 555 | }() |
| 556 | |
| 557 | select { |
| 558 | case err := <-serverErrors: |
| 559 | if err != nil && !errors.Is(err, http.ErrServerClosed) { |
| 560 | slog.Error("server error", "error", err) |
| 561 | } |
| 562 | case sig := <-stop: |
| 563 | slog.Info("shutdown initiated", "signal", sig.String()) |
| 564 | shutdownCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) |
| 565 | defer cancel() |
| 566 | |
| 567 | // Shutdown meter provider first |
| 568 | if meterProvider != nil { |
| 569 | slog.Info("shutting down OpenTelemetry meter provider") |
| 570 | if err := meterProvider.Shutdown(shutdownCtx); err != nil { |
| 571 | slog.Error("meter provider shutdown failed", "error", err) |
| 572 | } |
| 573 | } |
| 574 | |
| 575 | // Shutdown HTTP server |
| 576 | slog.Info("shutting down HTTP server") |
| 577 | if err := server.Shutdown(shutdownCtx); err != nil { |
| 578 | slog.Error("server shutdown failed", "error", err) |
| 579 | } else { |
| 580 | slog.Info("server shutdown completed gracefully") |
| 581 | } |
| 582 | |
| 583 | // Close deduplication log file if open |
| 584 | if dedupLogger != nil { |
| 585 | slog.Info("closing deduplication log file") |
| 586 | if err := dedupLogger.Close(); err != nil { |
| 587 | slog.Error("failed to close deduplication log file", "error", err) |
| 588 | } |
| 589 | } |
| 590 | } |
| 591 | |
| 592 | slog.Info("server exiting") |
| 593 | } |
| 594 | |
| 595 | // --- Helper Functions --- |
| 596 | func limitString(s string, maxLen int) string { |
| 597 | if len(s) <= maxLen { |
| 598 | return s |
| 599 | } |
| 600 | return s[:maxLen] + "..." |
| 601 | } |