diff --git a/README.md b/README.md index 7c34de6..f96f224 100644 --- a/README.md +++ b/README.md @@ -142,6 +142,151 @@ The library handles DNS resolution differently depending on the connection type: **Important for Privacy**: When using `socks5h://` or other remote DNS proxies, your local DNS servers will not see any queries for the target domains, maintaining better privacy and anonymity. + +### Metrics and Observability + +`gowarc` provides a `StatsRegistry` interface that allows you to integrate your own metrics collection system (Prometheus, Datadog, etc.). The library tracks various metrics including data written, deduplication statistics, and proxy usage. + +#### Using the StatsRegistry Interface + +The `StatsRegistry` interface can be found in [`stats.go`](stats.go). To implement your own metrics collection: + +```go +// Implement the StatsRegistry interface +type MyPrometheusRegistry struct { + // Your Prometheus registry fields +} + +func (r *MyPrometheusRegistry) RegisterCounter(name, help string, labelNames []string) warc.Counter { + // Return a Counter that wraps your Prometheus counter + // The Counter interface requires WithLabels() method for dimensional metrics +} + +func (r *MyPrometheusRegistry) RegisterGauge(name, help string, labelNames []string) warc.Gauge { + // Return a Gauge that wraps your Prometheus gauge +} + +func (r *MyPrometheusRegistry) RegisterHistogram(name, help string, buckets []int64, labelNames []string) warc.Histogram { + // Return a Histogram that wraps your Prometheus histogram +} + +// Pass your registry to the HTTP client +clientSettings := warc.HTTPClientSettings{ + StatsRegistry: &MyPrometheusRegistry{}, + // ... other settings +} +``` + +#### Available Metrics + +The library tracks the following metrics: + +- **`total_data_written`**: Total bytes written to WARC files +- **`local_deduped_bytes_total`**: Bytes saved through local deduplication +- **`local_deduped_total`**: Number of records deduplicated locally +- **`doppelganger_deduped_bytes_total`**: Bytes saved through Doppelganger deduplication +- **`doppelganger_deduped_total`**: Number of records deduplicated via Doppelganger +- **`cdx_deduped_bytes_total`**: Bytes saved through CDX deduplication +- **`cdx_deduped_total`**: Number of records deduplicated via CDX +- **`proxy_requests_total`**: Total requests through each proxy (with `proxy` label) +- **`proxy_errors_total`**: Total errors for each proxy (with `proxy` label) +- **`proxy_last_used_nanoseconds`**: Last usage timestamp for each proxy (with `proxy` label) + +#### Label Support + +Metrics support Prometheus-style labels for dimensional data: + +```go +// Register a counter with label dimensions +counter := registry.RegisterCounter("http_requests_total", "Total HTTP requests", []string{"method", "status"}) + +// Record metrics with specific label values +counter.WithLabels(warc.Labels{"method": "GET", "status": "200"}).Inc() +counter.WithLabels(warc.Labels{"method": "POST", "status": "201"}).Add(5) + +// Each unique label combination creates a separate metric series +``` + +**Interface Details**: See the complete interface contract in [`stats.go`](stats.go) for full implementation requirements. + +### Logging + +`gowarc` provides a `LogBackend` interface that allows you to integrate your logging solution (slog, zap, logrus, etc.). The library logs key events including connection establishment, DNS resolution, proxy selection, TLS handshakes, WARC file operations, and errors. + +#### Using the LogBackend Interface + +The `LogBackend` interface can be found in [`logging.go`](logging.go). The interface matches `slog.Logger` method signatures for easy integration: + +```go +// Example: Wrapping slog.Logger to implement LogBackend +type SlogAdapter struct { + logger *slog.Logger +} + +func (s *SlogAdapter) Debug(msg string, args ...any) { + s.logger.Debug(msg, args...) +} + +func (s *SlogAdapter) Info(msg string, args ...any) { + s.logger.Info(msg, args...) +} + +func (s *SlogAdapter) Warn(msg string, args ...any) { + s.logger.Warn(msg, args...) +} + +func (s *SlogAdapter) Error(msg string, args ...any) { + s.logger.Error(msg, args...) +} + +func (s *SlogAdapter) Log(ctx context.Context, level slog.Level, msg string, args ...any) { + s.logger.Log(ctx, level, msg, args...) +} + +// Configure with your logger +handler := slog.NewJSONHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelInfo}) +logger := slog.New(handler) + +clientSettings := warc.HTTPClientSettings{ + LogBackend: &SlogAdapter{logger: logger}, + // ... other settings +} +``` + +#### Log Events + +The library logs structured events with contextual key-value pairs: + +**Connection & Network:** +- Proxy selection and connection status +- Direct connection establishment +- DNS resolution results and failures +- TLS handshake success and failures + +**WARC Operations:** +- WARC record writing and file rotation +- Data written and compression events +- File creation and closure + +**Errors:** +- Connection failures (proxy and direct) +- DNS resolution errors +- TLS handshake failures +- WARC write errors + +#### Log Levels + +- **Debug**: Verbose operational details (DNS lookups, successful connections, record writes) +- **Info**: Important state changes (file rotation, new WARC files) +- **Warn**: Recoverable issues and fallbacks +- **Error**: Failures and exceptions (connection errors, DNS failures, write errors) + +Users control which log levels are recorded by configuring their logger implementation's level threshold. + +**Note**: The `LogBackend` interface is intended to eventually replace the `ErrChan` error reporting mechanism. For now, both are maintained for backward compatibility. + +**Interface Details**: See the complete interface contract in [`logging.go`](logging.go) for full implementation requirements. + ## CLI Tools In addition to the Go library, gowarc provides several command-line utilities for working with WARC files: diff --git a/client.go b/client.go index 9591e55..8c60275 100644 --- a/client.go +++ b/client.go @@ -4,7 +4,6 @@ import ( "net/http" "os" "sync" - "sync/atomic" "time" ) @@ -13,9 +12,52 @@ type Error struct { Func string } +// ProxyNetwork defines the network layer (IPv4/IPv6) a proxy can support +type ProxyNetwork int + +const ( + // ProxyNetworkUnset is the zero value and must not be used - forces explicit selection + ProxyNetworkUnset ProxyNetwork = iota + // ProxyNetworkAny means the proxy can be used for both IPv4 and IPv6 connections + ProxyNetworkAny + // ProxyNetworkIPv4 means the proxy should only be used for IPv4 connections + ProxyNetworkIPv4 + // ProxyNetworkIPv6 means the proxy should only be used for IPv6 connections + ProxyNetworkIPv6 +) + +// ProxyType defines the infrastructure type of a proxy +type ProxyType int + +const ( + // ProxyTypeAny means the proxy can be used for any type of request + ProxyTypeAny ProxyType = iota + // ProxyTypeMobile means the proxy uses mobile network infrastructure + ProxyTypeMobile + // ProxyTypeResidential means the proxy uses residential IP addresses + ProxyTypeResidential + // ProxyTypeDatacenter means the proxy uses datacenter infrastructure + ProxyTypeDatacenter +) + +// ProxyConfig defines the configuration for a single proxy +type ProxyConfig struct { + // URL is the proxy URL (e.g., "socks5://proxy.example.com:1080") + URL string + // Network specifies if this proxy supports IPv4, IPv6, or both + Network ProxyNetwork + // Type specifies the infrastructure type (Mobile, Residential, Datacenter, or Any) + Type ProxyType + // AllowedDomains is a list of glob patterns for domains this proxy should handle + // Examples: "*.example.com", "api.*.org" + // If empty, the proxy can be used for any domain + AllowedDomains []string +} + type HTTPClientSettings struct { RotatorSettings *RotatorSettings - Proxy string + Proxies []ProxyConfig + AllowDirectFallback bool TempDir string DiscardHook DiscardHook DNSServers []string @@ -39,13 +81,14 @@ type HTTPClientSettings struct { DisableIPv6 bool IPv6AnyIP bool DigestAlgorithm DigestAlgorithm + StatsRegistry StatsRegistry + LogBackend LogBackend } type CustomHTTPClient struct { interfacesWatcherStop chan bool WaitGroup *WaitGroupWithCount dedupeHashTable *sync.Map - ErrChan chan *Error WARCWriter chan *RecordBatch interfacesWatcherStarted chan bool http.Client @@ -64,15 +107,9 @@ type CustomHTTPClient struct { // If set to <= 0, the default value is DefaultMaxRAMUsageFraction. MaxRAMUsageFraction float64 randomLocalIP bool - DataTotal *atomic.Int64 - - CDXDedupeTotalBytes *atomic.Int64 - DoppelgangerDedupeTotalBytes *atomic.Int64 - LocalDedupeTotalBytes *atomic.Int64 - CDXDedupeTotal *atomic.Int64 - DoppelgangerDedupeTotal *atomic.Int64 - LocalDedupeTotal *atomic.Int64 + statsRegistry StatsRegistry + logBackend LogBackend } func (c *CustomHTTPClient) Close() error { @@ -91,7 +128,6 @@ func (c *CustomHTTPClient) Close() error { } wg.Wait() - close(c.ErrChan) if c.randomLocalIP { c.interfacesWatcherStop <- true @@ -106,16 +142,24 @@ func (c *CustomHTTPClient) Close() error { func NewWARCWritingHTTPClient(HTTPClientSettings HTTPClientSettings) (httpClient *CustomHTTPClient, err error) { httpClient = new(CustomHTTPClient) - // Initialize counters - httpClient.DataTotal = &DataTotal - - httpClient.CDXDedupeTotalBytes = &CDXDedupeTotalBytes - httpClient.DoppelgangerDedupeTotalBytes = &DoppelgangerDedupeTotalBytes - httpClient.LocalDedupeTotalBytes = &LocalDedupeTotalBytes + // Initialize stats registry + if HTTPClientSettings.StatsRegistry != nil { + httpClient.statsRegistry = HTTPClientSettings.StatsRegistry + HTTPClientSettings.RotatorSettings.StatsRegistry = HTTPClientSettings.StatsRegistry + } else { + localStatsRegistry := newLocalRegistry() + httpClient.statsRegistry = localStatsRegistry + HTTPClientSettings.RotatorSettings.StatsRegistry = localStatsRegistry + } - httpClient.CDXDedupeTotal = &CDXDedupeTotal - httpClient.DoppelgangerDedupeTotal = &DoppelgangerDedupeTotal - httpClient.LocalDedupeTotal = &LocalDedupeTotal + // Initialize log backend + if HTTPClientSettings.LogBackend != nil { + httpClient.logBackend = HTTPClientSettings.LogBackend + HTTPClientSettings.RotatorSettings.LogBackend = HTTPClientSettings.LogBackend + } else { + httpClient.logBackend = &noopLogger{} + HTTPClientSettings.RotatorSettings.LogBackend = &noopLogger{} + } // Configure random local IP httpClient.randomLocalIP = HTTPClientSettings.RandomLocalIP @@ -142,9 +186,6 @@ func NewWARCWritingHTTPClient(HTTPClientSettings HTTPClientSettings) (httpClient // Set a hook to determine if we should discard a response httpClient.DiscardHook = HTTPClientSettings.DiscardHook - // Create an error channel for sending WARC errors through - httpClient.ErrChan = make(chan *Error) - // Toggle verification of certificates // InsecureSkipVerify expects the opposite of the verifyCerts flag, as such we flip it. httpClient.verifyCerts = !HTTPClientSettings.VerifyCerts @@ -216,7 +257,7 @@ func NewWARCWritingHTTPClient(HTTPClientSettings HTTPClientSettings) (httpClient httpClient.ConnReadDeadline = HTTPClientSettings.ConnReadDeadline // Configure custom dialer / transport - customDialer, err := newCustomDialer(httpClient, HTTPClientSettings.Proxy, HTTPClientSettings.DialTimeout, HTTPClientSettings.DNSRecordsTTL, HTTPClientSettings.DNSResolutionTimeout, HTTPClientSettings.DNSCacheSize, HTTPClientSettings.DNSServers, HTTPClientSettings.DNSConcurrency, HTTPClientSettings.DisableIPv4, HTTPClientSettings.DisableIPv6) + customDialer, err := newCustomDialer(httpClient, HTTPClientSettings.Proxies, HTTPClientSettings.AllowDirectFallback, HTTPClientSettings.DialTimeout, HTTPClientSettings.DNSRecordsTTL, HTTPClientSettings.DNSResolutionTimeout, HTTPClientSettings.DNSCacheSize, HTTPClientSettings.DNSServers, HTTPClientSettings.DNSConcurrency, HTTPClientSettings.DisableIPv4, HTTPClientSettings.DisableIPv6) if err != nil { return nil, err } diff --git a/client_test.go b/client_test.go index 99ef81e..560100d 100644 --- a/client_test.go +++ b/client_test.go @@ -10,6 +10,7 @@ import ( "errors" "fmt" "io" + "log/slog" "math/big" "net" "net/http" @@ -39,6 +40,8 @@ func defaultRotatorSettings(t *testing.T) *RotatorSettings { err error ) + rotatorSettings.StatsRegistry = &localRegistry{} + rotatorSettings.LogBackend = &noopLogger{} rotatorSettings.Prefix = "TEST" rotatorSettings.OutputDirectory, err = os.MkdirTemp("", "warc-tests-") if err != nil { @@ -105,19 +108,6 @@ func sumRecordContentLengths(path string) (int64, error) { return total, nil } -// Helper function used in many tests -func drainErrChan(t *testing.T, errChan chan *Error) func() { - var wg sync.WaitGroup - wg.Add(1) - go func() { - defer wg.Done() - for err := range errChan { - t.Errorf("Error writing to WARC: %s", err.Err.Error()) - } - }() - return func() { wg.Wait() } -} - func newTestImageServer(t testing.TB, st int) *httptest.Server { return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { fileBytes, err := os.ReadFile(path.Join("testdata", "image.svg")) @@ -160,9 +150,6 @@ func TestHTTPClient(t *testing.T) { err error ) - // Reset counter to 0 - DataTotal.Store(0) - // init test HTTP endpoint server := newTestImageServer(t, http.StatusOK) defer server.Close() @@ -172,7 +159,6 @@ func TestHTTPClient(t *testing.T) { if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) req, err := http.NewRequest("GET", server.URL+"/testdata/image.svg", nil) if err != nil { @@ -188,7 +174,6 @@ func TestHTTPClient(t *testing.T) { io.Copy(io.Discard, resp.Body) httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -207,7 +192,7 @@ func TestHTTPClient(t *testing.T) { } // verify that the remote dedupe count is correct - dataTotal := httpClient.DataTotal.Load() + dataTotal := httpClient.statsRegistry.RegisterCounter(totalDataWritten, totalDataWrittenHelp, nil).WithLabels(nil).Get() if dataTotal != expectedPayloadBytes { t.Fatalf("total bytes downloaded mismatch, expected %d got %d", expectedPayloadBytes, dataTotal) } @@ -223,22 +208,18 @@ func TestHTTPClientRequestFailing(t *testing.T) { server := newTestImageServer(t, http.StatusOK) defer server.Close() + // init test logger to capture errors + testLog := NewTestLogger() + // init the HTTP client responsible for recording HTTP(s) requests / responses - httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{RotatorSettings: rotatorSettings}) + httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ + RotatorSettings: rotatorSettings, + LogBackend: testLog, + }) if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - errCh := make(chan *Error, 1) - var errChWg sync.WaitGroup - errChWg.Add(1) - go func() { - defer errChWg.Done() - for err := range httpClient.ErrChan { - errCh <- err - } - }() - // Prepare some dummy data and configure our error injector data := []byte("this is some test data") erc := &errorReadCloser{data: data, readBefore: 10} // allow 10 bytes before error @@ -249,33 +230,21 @@ func TestHTTPClientRequestFailing(t *testing.T) { } _, err = httpClient.Do(req) - if err == nil { - select { - case recv := <-errCh: - if recv == nil { - t.Fatal("expected error via ErrChan but channel closed without value") - } - case <-time.After(2 * time.Second): - t.Fatal("expected error on Do or via ErrChan, got none") - } - } else { - t.Logf("got expected error: %v", err) + if err == nil || !strings.Contains(err.Error(), "injected read error") { + t.Fatal("expected \"injected read error\" error, got nil") } httpClient.Close() - errChWg.Wait() - close(errCh) } func TestHTTPClientConnReadDeadline(t *testing.T) { var ( rotatorSettings = defaultRotatorSettings(t) - errWg sync.WaitGroup err error ) // 1) Set up a test server that sends its response slowly, in chunks - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "text/plain") w.WriteHeader(http.StatusOK) @@ -301,14 +270,6 @@ func TestHTTPClientConnReadDeadline(t *testing.T) { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - // Read any WARC-writing errors - errWg.Add(1) - go func() { - defer errWg.Done() - for range httpClient.ErrChan { - } - }() - // 3) Create a request req, err := http.NewRequest("GET", server.URL, nil) if err != nil { @@ -335,7 +296,6 @@ func TestHTTPClientConnReadDeadline(t *testing.T) { func TestHTTPClientContextCancellation(t *testing.T) { var ( rotatorSettings = defaultRotatorSettings(t) - errWg sync.WaitGroup err error ) @@ -363,15 +323,6 @@ func TestHTTPClientContextCancellation(t *testing.T) { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - // Read any WARC-writing errors - errWg.Add(1) - go func() { - defer errWg.Done() - for _ = range httpClient.ErrChan { - // t.Errorf("Error writing to WARC: %s", e.Err.Error()) - } - }() - // 3) Create a request with a cancellable context ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -424,11 +375,12 @@ func TestHTTPClientWithFeedbackChan(t *testing.T) { defer server.Close() // init the HTTP client responsible for recording HTTP(s) requests / responses - httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{RotatorSettings: rotatorSettings}) + httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ + RotatorSettings: rotatorSettings, + }) if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) req, err := http.NewRequest("GET", server.URL+"/testdata/image.svg", nil) if err != nil { @@ -449,7 +401,6 @@ func TestHTTPClientWithFeedbackChan(t *testing.T) { <-feedbackCh httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -499,6 +450,9 @@ func TestHTTPClientTLSHandshakeTimeout(t *testing.T) { serverURL := "https://" + ln.Addr().String() + // 4) Set up a test logger to capture errors + testLog := NewTestLogger() + // 5) Create the WARC-writing HTTP client // The critical part here is enforcing the handshake timeout. // (Exact field names may differ based on your library.) @@ -506,11 +460,11 @@ func TestHTTPClientTLSHandshakeTimeout(t *testing.T) { RotatorSettings: rotatorSettings, TLSHandshakeTimeout: 1 * time.Second, // <--- The key line VerifyCerts: true, // or "VerifyCerts: false" depending on your lib + LogBackend: testLog, }) if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %v", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) // 6) Attempt the GET, which should fail due to TLS handshake delay req, err := http.NewRequest("GET", serverURL, nil) @@ -529,8 +483,25 @@ func TestHTTPClientTLSHandshakeTimeout(t *testing.T) { t.Logf("Got expected error: %v", err) } + // Check if an error is logged with the expected message + gotErr := false + logEntries := make([]LogEntry, 0) + if testLog.HasLevel(slog.LevelError) { + for _, entry := range testLog.Entries() { + if entry.Level == slog.LevelError && strings.Contains(entry.Message, "TLS handshake failed") && strings.Contains(entry.Args[3].(error).Error(), "context deadline exceeded") { + gotErr = true + } else if entry.Level == slog.LevelError { + t.Logf("Unexpected log entry: %v", entry) + } + logEntries = append(logEntries, entry) + } + } + + if !gotErr { + t.Fatalf("Expected TLS handshake failed error in logs, got %v", logEntries) + } + httpClient.Close() - waitForErrors() <-doneChan // Wait for the server goroutine to exit } @@ -538,7 +509,6 @@ func TestHTTPClientTLSHandshakeTimeout(t *testing.T) { func TestHTTPClientServerClosingConnection(t *testing.T) { var ( rotatorSettings = defaultRotatorSettings(t) - errWg sync.WaitGroup err error ) @@ -573,19 +543,13 @@ func TestHTTPClientServerClosingConnection(t *testing.T) { defer server.Close() // init the HTTP client responsible for recording HTTP(s) requests / responses - httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{RotatorSettings: rotatorSettings}) + httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ + RotatorSettings: rotatorSettings, + }) if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - errWg.Add(1) - go func() { - defer errWg.Done() - for _ = range httpClient.ErrChan { - // We expect an error here, so we don't need to log it - } - }() - req, err := http.NewRequest("GET", server.URL, nil) if err != nil { t.Fatal(err) @@ -618,14 +582,16 @@ func TestHTTPClientDNSFailure(t *testing.T) { err error ) + testLog := NewTestLogger() + // Initialize the WARC-writing HTTP client httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ RotatorSettings: rotatorSettings, + LogBackend: testLog, }) if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) // Use a guaranteed-nonresolvable domain req, err := http.NewRequest("GET", "http://should-not-resolve.example.invalid", nil) @@ -644,8 +610,23 @@ func TestHTTPClientDNSFailure(t *testing.T) { t.Logf("Got expected DNS error: %v", err) } + // Verify the error was logged + gotErr := false + logEntries := make([]LogEntry, 0) + if testLog.HasLevel(slog.LevelError) { + for _, entry := range testLog.Entries() { + if entry.Level == slog.LevelError && strings.Contains(entry.Message, "DNS resolution failed") && strings.Contains(entry.Args[3].(error).Error(), "failed to resolve DNS: A error: no TYPE=A record found, AAAA error: no TYPE=AAAA record found") { + gotErr = true + } + logEntries = append(logEntries, entry) + } + } + + if !gotErr { + t.Fatalf("Expected DNS resolution failed error in logs, got %v", logEntries) + } + httpClient.Close() - waitForErrors() } func TestHTTPClientWithProxy(t *testing.T) { @@ -688,11 +669,16 @@ func TestHTTPClientWithProxy(t *testing.T) { // init the HTTP client responsible for recording HTTP(s) requests / responses httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ RotatorSettings: rotatorSettings, - Proxy: fmt.Sprintf("socks5://%s", proxyAddr)}) + Proxies: []ProxyConfig{ + { + URL: fmt.Sprintf("socks5://%s", proxyAddr), + Network: ProxyNetworkAny, + Type: ProxyTypeAny, + }, + }}) if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) req, err := http.NewRequest("GET", server.URL, nil) if err != nil { @@ -708,7 +694,6 @@ func TestHTTPClientWithProxy(t *testing.T) { io.Copy(io.Discard, resp.Body) httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -732,11 +717,12 @@ func TestHTTPClientConcurrent(t *testing.T) { defer server.Close() // init the HTTP client responsible for recording HTTP(s) requests / responses - httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{RotatorSettings: rotatorSettings}) + httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ + RotatorSettings: rotatorSettings, + }) if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) wg.Add(concurrency) for i := 0; i < concurrency; i++ { @@ -746,13 +732,11 @@ func TestHTTPClientConcurrent(t *testing.T) { req, err := http.NewRequest("GET", server.URL, nil) req.Close = true if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } resp, err := httpClient.Do(req) if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } defer resp.Body.Close() @@ -765,7 +749,6 @@ func TestHTTPClientConcurrent(t *testing.T) { wg.Wait() httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -790,11 +773,12 @@ func TestHTTPClientMultiWARCWriters(t *testing.T) { defer server.Close() // init the HTTP client responsible for recording HTTP(s) requests / responses - httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{RotatorSettings: rotatorSettings}) + httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ + RotatorSettings: rotatorSettings, + }) if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) wg.Add(concurrency) for i := 0; i < concurrency; i++ { @@ -804,13 +788,11 @@ func TestHTTPClientMultiWARCWriters(t *testing.T) { req, err := http.NewRequest("GET", server.URL, nil) req.Close = true if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } resp, err := httpClient.Do(req) if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } defer resp.Body.Close() @@ -823,7 +805,6 @@ func TestHTTPClientMultiWARCWriters(t *testing.T) { wg.Wait() httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -846,10 +827,6 @@ func TestHTTPClientLocalDedupe(t *testing.T) { err error ) - // Reset counter to 0 - LocalDedupeTotal.Store(0) - LocalDedupeTotalBytes.Store(0) - // init test HTTP endpoint server := newTestImageServer(t, http.StatusOK) defer server.Close() @@ -864,7 +841,6 @@ func TestHTTPClientLocalDedupe(t *testing.T) { if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) for i := 0; i < 2; i++ { req, err := http.NewRequest("GET", server.URL, nil) @@ -884,7 +860,6 @@ func TestHTTPClientLocalDedupe(t *testing.T) { } httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -897,18 +872,13 @@ func TestHTTPClientLocalDedupe(t *testing.T) { } // verify that the local dedupe count is correct - if LocalDedupeTotalBytes.Load() != 26872 { - t.Fatalf("local dedupe total bytes mismatch, expected: 26872 got: %d", LocalDedupeTotalBytes.Load()) - } - - // Ensure that HTTP client results work correctly as well - if httpClient.LocalDedupeTotalBytes.Load() != 26872 { - t.Fatalf("local dedupe total bytes mismatch, expected: 26872 got: %d", httpClient.LocalDedupeTotalBytes.Load()) + if httpClient.statsRegistry.RegisterCounter(localDedupedBytesTotal, localDedupedBytesTotalHelp, nil).WithLabels(nil).Get() != 26872 { + t.Fatalf("local dedupe total bytes mismatch, expected: 26872 got: %d", httpClient.statsRegistry.RegisterCounter(localDedupedBytesTotal, localDedupedBytesTotalHelp, nil).WithLabels(nil).Get()) } // 1 is expected due to requiring one request to enter into the table. - if httpClient.LocalDedupeTotal.Load() != 1 { - t.Fatalf("local dedupe total mismatch, expected: 1 got: %d", httpClient.LocalDedupeTotal.Load()) + if httpClient.statsRegistry.RegisterCounter(localDedupedTotal, localDedupedTotalHelp, nil).WithLabels(nil).Get() != 1 { + t.Fatalf("local dedupe total mismatch, expected: 1 got: %d", httpClient.statsRegistry.RegisterCounter(localDedupedTotal, localDedupedTotalHelp, nil).WithLabels(nil).Get()) } } @@ -922,11 +892,7 @@ func TestHTTPClientRemoteDedupe(t *testing.T) { // init test HTTP endpoint mux := http.NewServeMux() - // Reset counter to 0 - CDXDedupeTotal.Store(0) - CDXDedupeTotalBytes.Store(0) - - mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + mux.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) { fileBytes, err := os.ReadFile(path.Join("testdata", "image.svg")) if err != nil { t.Fatal(err) @@ -937,7 +903,7 @@ func TestHTTPClientRemoteDedupe(t *testing.T) { w.Write(fileBytes) }) - mux.HandleFunc(dedupePath, func(w http.ResponseWriter, r *http.Request) { + mux.HandleFunc(dedupePath, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "text/plain;charset=UTF-8") w.WriteHeader(http.StatusOK) w.Write([]byte(dedupeResp)) @@ -958,7 +924,6 @@ func TestHTTPClientRemoteDedupe(t *testing.T) { if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) for i := 0; i < 4; i++ { req, err := http.NewRequest("GET", server.URL, nil) @@ -978,7 +943,6 @@ func TestHTTPClientRemoteDedupe(t *testing.T) { } httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -990,18 +954,13 @@ func TestHTTPClientRemoteDedupe(t *testing.T) { testFileRevisitVailidity(t, path, "2022-03-20T00:25:18Z", "sha1:UIRWL5DFIPQ4MX3D3GFHM2HCVU3TZ6I3", false) } - // verify that the remote dedupe count is correct - if CDXDedupeTotalBytes.Load() != 107488 { - t.Fatalf("remote dedupe total bytes mismatch, expected: 107488 got: %d", CDXDedupeTotalBytes.Load()) - } - - // Ensure that HTTP client results work correctly as well - if httpClient.CDXDedupeTotalBytes.Load() != 107488 { - t.Fatalf("remote dedupe total bytes mismatch, expected: 107488 got: %d", httpClient.CDXDedupeTotalBytes.Load()) + // verify that the CDX dedupe count is correct + if httpClient.statsRegistry.RegisterCounter(cdxDedupedBytesTotal, cdxDedupedBytesTotalHelp, nil).WithLabels(nil).Get() != 107488 { + t.Fatalf("CDX dedupe total bytes mismatch, expected: 26872 got: %d", httpClient.statsRegistry.RegisterCounter(cdxDedupedBytesTotal, cdxDedupedBytesTotalHelp, nil).WithLabels(nil).Get()) } - if httpClient.CDXDedupeTotal.Load() != 4 { - t.Fatalf("remote dedupe total mismatch, expected: 4 got: %d", httpClient.CDXDedupeTotal.Load()) + if httpClient.statsRegistry.RegisterCounter(cdxDedupedTotal, cdxDedupedTotalHelp, nil).WithLabels(nil).Get() != 4 { + t.Fatalf("CDX dedupe total mismatch, expected: 1 got: %d", httpClient.statsRegistry.RegisterCounter(cdxDedupedTotal, cdxDedupedTotalHelp, nil).WithLabels(nil).Get()) } } @@ -1010,17 +969,12 @@ func TestHTTPClientDoppelgangerDedupe(t *testing.T) { dedupePath = "/api/records/UIRWL5DFIPQ4MX3D3GFHM2HCVU3TZ6I3" dedupeResp = "{\"id\":\"UIRWL5DFIPQ4MX3D3GFHM2HCVU3TZ6I3\",\"uri\":\"https://upload.wikimedia.org/wikipedia/commons/5/55/Blason_ville_fr_Sarlat-la-Can%C3%A9da_%28Dordogne%29.svg\",\"date\":20220320002518}" rotatorSettings = defaultRotatorSettings(t) - errWg sync.WaitGroup err error ) // init test HTTP endpoint mux := http.NewServeMux() - // Reset counter to 0 - DoppelgangerDedupeTotal.Store(0) - DoppelgangerDedupeTotalBytes.Store(0) - - mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + mux.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) { fileBytes, err := os.ReadFile(path.Join("testdata", "image.svg")) if err != nil { t.Fatal(err) @@ -1031,7 +985,7 @@ func TestHTTPClientDoppelgangerDedupe(t *testing.T) { w.Write(fileBytes) }) - mux.HandleFunc(dedupePath, func(w http.ResponseWriter, r *http.Request) { + mux.HandleFunc(dedupePath, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) w.Write([]byte(dedupeResp)) @@ -1054,14 +1008,6 @@ func TestHTTPClientDoppelgangerDedupe(t *testing.T) { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - errWg.Add(1) - go func() { - defer errWg.Done() - for err := range httpClient.ErrChan { - t.Errorf("Error writing to WARC: %s", err.Err.Error()) - } - }() - for i := 0; i < 4; i++ { req, err := http.NewRequest("GET", server.URL, nil) if err != nil { @@ -1089,18 +1035,13 @@ func TestHTTPClientDoppelgangerDedupe(t *testing.T) { testFileRevisitVailidity(t, path, "2022-03-20T00:25:18Z", "sha1:UIRWL5DFIPQ4MX3D3GFHM2HCVU3TZ6I3", false) } - // verify that the remote dedupe count is correct - if DoppelgangerDedupeTotalBytes.Load() != 107488 { - t.Fatalf("remote dedupe total bytes mismatch, expected: 107488 got: %d", DoppelgangerDedupeTotalBytes.Load()) + // verify that the Doppelganger count is correct + if httpClient.statsRegistry.RegisterCounter(doppelgangerDedupedBytesTotal, doppelgangerDedupedBytesTotalHelp, nil).WithLabels(nil).Get() != 107488 { + t.Fatalf("Doppelganger total bytes mismatch, expected: 26872 got: %d", httpClient.statsRegistry.RegisterCounter(doppelgangerDedupedBytesTotal, doppelgangerDedupedBytesTotalHelp, nil).WithLabels(nil).Get()) } - // Ensure that HTTP client results work correctly as well - if httpClient.DoppelgangerDedupeTotalBytes.Load() != 107488 { - t.Fatalf("remote dedupe total bytes mismatch, expected: 107488 got: %d", httpClient.DoppelgangerDedupeTotalBytes.Load()) - } - - if httpClient.DoppelgangerDedupeTotal.Load() != 4 { - t.Fatalf("remote dedupe total mismatch, expected: 4 got: %d", httpClient.DoppelgangerDedupeTotal.Load()) + if httpClient.statsRegistry.RegisterCounter(doppelgangerDedupedTotal, doppelgangerDedupedTotalHelp, nil).WithLabels(nil).Get() != 4 { + t.Fatalf("Doppelganger total mismatch, expected: 1 got: %d", httpClient.statsRegistry.RegisterCounter(doppelgangerDedupedTotal, doppelgangerDedupedTotalHelp, nil).WithLabels(nil).Get()) } } @@ -1110,10 +1051,6 @@ func TestHTTPClientDedupeEmptyPayload(t *testing.T) { err error ) - // Reset counter to 0 - LocalDedupeTotal.Store(0) - LocalDedupeTotalBytes.Store(0) - // init test HTTP endpoint server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // Empty. This is intentional to mirror 3I42H3S6NNFQ2MSVX7XZKYAYSCX5QBYJ. @@ -1134,7 +1071,6 @@ func TestHTTPClientDedupeEmptyPayload(t *testing.T) { if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) for i := 0; i < 2; i++ { req, err := http.NewRequest("GET", server.URL, nil) @@ -1154,7 +1090,6 @@ func TestHTTPClientDedupeEmptyPayload(t *testing.T) { } httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -1167,24 +1102,18 @@ func TestHTTPClientDedupeEmptyPayload(t *testing.T) { } // verify that the local dedupe count is correct - if LocalDedupeTotalBytes.Load() != 0 { - t.Fatalf("local dedupe total bytes mismatch, expected: 0 got: %d", LocalDedupeTotalBytes.Load()) - } - - // Ensure that HTTP client results work correctly as well - if httpClient.LocalDedupeTotalBytes.Load() != 0 { - t.Fatalf("local dedupe total bytes mismatch, expected: 0 got: %d", httpClient.LocalDedupeTotalBytes.Load()) + if httpClient.statsRegistry.RegisterCounter(localDedupedBytesTotal, localDedupedBytesTotalHelp, nil).WithLabels(nil).Get() != 0 { + t.Fatalf("local dedupe total bytes mismatch, expected: 26872 got: %d", httpClient.statsRegistry.RegisterCounter(localDedupedBytesTotal, localDedupedBytesTotalHelp, nil).WithLabels(nil).Get()) } - if httpClient.LocalDedupeTotal.Load() != 0 { - t.Fatalf("local dedupe total mismatch, expected: 0 got: %d", httpClient.LocalDedupeTotal.Load()) + if httpClient.statsRegistry.RegisterCounter(localDedupedTotal, localDedupedTotalHelp, nil).WithLabels(nil).Get() != 0 { + t.Fatalf("local dedupe total mismatch, expected: 1 got: %d", httpClient.statsRegistry.RegisterCounter(localDedupedTotal, localDedupedTotalHelp, nil).WithLabels(nil).Get()) } } func TestHTTPClientDiscardHook(t *testing.T) { var ( rotatorSettings = defaultRotatorSettings(t) - errWg sync.WaitGroup err error ) @@ -1194,9 +1123,12 @@ func TestHTTPClientDiscardHook(t *testing.T) { server := newTestImageServer(t, http.StatusTooManyRequests) defer server.Close() + testLog := NewTestLogger() + // init the HTTP client responsible for recording HTTP(s) requests / responses httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ RotatorSettings: rotatorSettings, + LogBackend: testLog, // Set up a discard hook to discard 429 responses DiscardHook: func(resp *http.Response) (bool, string) { if resp.StatusCode != http.StatusTooManyRequests { @@ -1210,25 +1142,6 @@ func TestHTTPClientDiscardHook(t *testing.T) { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - errWg.Add(1) - go func() { - defer errWg.Done() - for err := range httpClient.ErrChan { - // validate 429 filtering as well as error reporting by url - discardErr, ok := err.Err.(*DiscardHookError) - if !ok { - t.Errorf("Expected DiscardHookError, got: %T, error: %v", err.Err, err) - continue - } - if discardErr.URL != server.URL+"/" { - t.Errorf("Expected URL %s, got: %s", server.URL+"/", discardErr.URL) - } - if discardErr.Reason != expectedReason { - t.Errorf("Expected Reason %s, got: %s", expectedReason, discardErr.Reason) - } - } - }() - req, err := http.NewRequest("GET", server.URL, nil) if err != nil { t.Fatal(err) @@ -1244,6 +1157,22 @@ func TestHTTPClientDiscardHook(t *testing.T) { httpClient.Close() + // Verify the error was logged + gotErr := false + logEntries := make([]LogEntry, 0) + if testLog.HasLevel(slog.LevelError) { + for _, entry := range testLog.Entries() { + if entry.Level == slog.LevelError && strings.Contains(entry.Message, "error reading from connection") && strings.Contains(entry.Args[3].(*DiscardHookError).Error(), "response was blocked by DiscardHook") && strings.Contains(entry.Args[3].(*DiscardHookError).Error(), "429 response") { + gotErr = true + } + logEntries = append(logEntries, entry) + } + } + + if !gotErr { + t.Fatalf("Expected \"error reading from connection\" error in logs, got %v", logEntries) + } + files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { t.Fatal(err) @@ -1279,7 +1208,6 @@ func TestHTTPClientPayloadLargerThan2MB(t *testing.T) { if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) req, err := http.NewRequest("GET", server.URL, nil) if err != nil { @@ -1295,7 +1223,6 @@ func TestHTTPClientPayloadLargerThan2MB(t *testing.T) { io.Copy(io.Discard, resp.Body) httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -1334,7 +1261,6 @@ func TestConcurrentHTTPClientPayloadLargerThan2MB(t *testing.T) { if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) wg.Add(concurrency) for i := 0; i < concurrency; i++ { @@ -1344,13 +1270,11 @@ func TestConcurrentHTTPClientPayloadLargerThan2MB(t *testing.T) { req, err := http.NewRequest("GET", server.URL, nil) req.Close = true if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } resp, err := httpClient.Do(req) if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } @@ -1363,7 +1287,6 @@ func TestConcurrentHTTPClientPayloadLargerThan2MB(t *testing.T) { wg.Wait() httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -1395,7 +1318,6 @@ func TestHTTPClientWithSelfSignedCertificate(t *testing.T) { if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) req, err := http.NewRequest("GET", server.URL, nil) if err != nil { @@ -1411,7 +1333,6 @@ func TestHTTPClientWithSelfSignedCertificate(t *testing.T) { io.Copy(io.Discard, resp.Body) httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -1448,7 +1369,6 @@ func TestWARCWritingWithDisallowedCertificate(t *testing.T) { if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) req, err := http.NewRequest("GET", server.URL, nil) if err != nil { @@ -1467,7 +1387,6 @@ func TestWARCWritingWithDisallowedCertificate(t *testing.T) { } httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -1495,7 +1414,6 @@ func TestHTTPClientFullOnDisk(t *testing.T) { if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) req, err := http.NewRequest("GET", server.URL, nil) if err != nil { @@ -1511,7 +1429,6 @@ func TestHTTPClientFullOnDisk(t *testing.T) { io.Copy(io.Discard, resp.Body) httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -1526,7 +1443,6 @@ func TestHTTPClientFullOnDisk(t *testing.T) { func TestHTTPClientWithoutIoCopy(t *testing.T) { var ( rotatorSettings = defaultRotatorSettings(t) - errWg sync.WaitGroup err error ) @@ -1537,23 +1453,17 @@ func TestHTTPClientWithoutIoCopy(t *testing.T) { server := newTestImageServer(t, http.StatusOK) defer server.Close() + testLog := NewTestLogger() + // init the HTTP client responsible for recording HTTP(s) requests / responses - httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{RotatorSettings: rotatorSettings}) + httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ + RotatorSettings: rotatorSettings, + LogBackend: testLog, + }) if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - errWg.Add(1) - go func() { - defer errWg.Done() - for err := range httpClient.ErrChan { - // validate 429 filtering as well as error reporting by url - if strings.Contains(err.Err.Error(), "SHA1 ran into an unrecoverable error url") { - t.Errorf("Error writing to WARC: %s", err.Err.Error()) - } - } - }() - req, err := http.NewRequest("GET", server.URL, nil) if err != nil { t.Fatal(err) @@ -1569,6 +1479,21 @@ func TestHTTPClientWithoutIoCopy(t *testing.T) { httpClient.Close() + gotErr := false + logEntries := make([]LogEntry, 0) + if testLog.HasLevel(slog.LevelError) { + for _, entry := range testLog.Entries() { + if entry.Level == slog.LevelError && strings.Contains(entry.Message, "error reading from connection") && strings.Contains(entry.Args[3].(error).Error(), "readResponse: payload digest calculation failed: unexpected EOF") { + gotErr = true + } + logEntries = append(logEntries, entry) + } + } + + if !gotErr { + t.Fatalf("Expected \"error reading from connection\" error in logs, got %v", logEntries) + } + files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { t.Fatal(err) @@ -1598,7 +1523,6 @@ func TestHTTPClientWithoutChunkEncoding(t *testing.T) { if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) req, err := http.NewRequest("GET", server.URL, nil) if err != nil { @@ -1614,7 +1538,6 @@ func TestHTTPClientWithoutChunkEncoding(t *testing.T) { io.Copy(io.Discard, resp.Body) httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -1642,7 +1565,6 @@ func TestHTTPClientWithZStandard(t *testing.T) { if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) req, err := http.NewRequest("GET", server.URL, nil) if err != nil { @@ -1658,7 +1580,6 @@ func TestHTTPClientWithZStandard(t *testing.T) { io.Copy(io.Discard, resp.Body) httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -1688,7 +1609,6 @@ func TestHTTPClientWithZStandardDictionary(t *testing.T) { if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - waitForErrors := drainErrChan(t, httpClient.ErrChan) req, err := http.NewRequest("GET", server.URL, nil) if err != nil { @@ -1704,7 +1624,6 @@ func TestHTTPClientWithZStandardDictionary(t *testing.T) { io.Copy(io.Discard, resp.Body) httpClient.Close() - waitForErrors() files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { @@ -1763,9 +1682,12 @@ func TestHTTPClientWithIPv4Disabled(t *testing.T) { rotatorSettings := defaultRotatorSettings(t) + testLog := NewTestLogger() + httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ RotatorSettings: rotatorSettings, DisableIPv4: true, + LogBackend: testLog, }) if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) @@ -1791,6 +1713,21 @@ func TestHTTPClientWithIPv4Disabled(t *testing.T) { httpClient.Close() + gotErr := false + logEntries := make([]LogEntry, 0) + if testLog.HasLevel(slog.LevelError) { + for _, entry := range testLog.Entries() { + if entry.Level == slog.LevelError && strings.Contains(entry.Message, "direct connection failed") && strings.Contains(entry.Args[3].(error).Error(), "dial tcp6: address 127.0.0.1: no suitable address found") { + gotErr = true + } + logEntries = append(logEntries, entry) + } + } + + if !gotErr { + t.Fatalf("Expected \"direct connection failed\" error in logs, got %v", logEntries) + } + files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { t.Fatal(err) @@ -1810,9 +1747,12 @@ func TestHTTPClientWithIPv6Disabled(t *testing.T) { rotatorSettings := defaultRotatorSettings(t) + testLog := NewTestLogger() + httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ RotatorSettings: rotatorSettings, DisableIPv6: true, + LogBackend: testLog, }) if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) @@ -1838,6 +1778,21 @@ func TestHTTPClientWithIPv6Disabled(t *testing.T) { httpClient.Close() + gotErr := false + logEntries := make([]LogEntry, 0) + if testLog.HasLevel(slog.LevelError) { + for _, entry := range testLog.Entries() { + if entry.Level == slog.LevelError && strings.Contains(entry.Message, "direct connection failed") && strings.Contains(entry.Args[3].(error).Error(), "dial tcp4: address ::1: no suitable address found") { + gotErr = true + } + logEntries = append(logEntries, entry) + } + } + + if !gotErr { + t.Fatalf("Expected \"direct connection failed\" error in logs, got %v", logEntries) + } + files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { t.Fatal(err) @@ -1853,7 +1808,6 @@ func BenchmarkConcurrentUnder2MB(b *testing.B) { var ( rotatorSettings = defaultBenchmarkRotatorSettings(b) wg sync.WaitGroup - errWg sync.WaitGroup err error ) @@ -1876,14 +1830,6 @@ func BenchmarkConcurrentUnder2MB(b *testing.B) { b.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - errWg.Add(1) - go func() { - defer errWg.Done() - for err := range httpClient.ErrChan { - b.Errorf("Error writing to WARC: %s", err.Err.Error()) - } - }() - wg.Add(b.N) for n := 0; n < b.N; n++ { go func() { @@ -1891,13 +1837,11 @@ func BenchmarkConcurrentUnder2MB(b *testing.B) { req, err := http.NewRequest("GET", server.URL, nil) if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } resp, err := httpClient.Do(req) if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } defer resp.Body.Close() @@ -1914,7 +1858,6 @@ func BenchmarkConcurrentUnder2MBZStandard(b *testing.B) { var ( rotatorSettings = defaultBenchmarkRotatorSettings(b) wg sync.WaitGroup - errWg sync.WaitGroup err error ) rotatorSettings.Compression = "ZSTD" @@ -1938,14 +1881,6 @@ func BenchmarkConcurrentUnder2MBZStandard(b *testing.B) { b.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - errWg.Add(1) - go func() { - defer errWg.Done() - for err := range httpClient.ErrChan { - b.Errorf("Error writing to WARC: %s", err.Err.Error()) - } - }() - wg.Add(b.N) for n := 0; n < b.N; n++ { go func() { @@ -1953,13 +1888,11 @@ func BenchmarkConcurrentUnder2MBZStandard(b *testing.B) { req, err := http.NewRequest("GET", server.URL, nil) if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } resp, err := httpClient.Do(req) if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } defer resp.Body.Close() @@ -1976,7 +1909,6 @@ func BenchmarkConcurrentOver2MB(b *testing.B) { var ( rotatorSettings = defaultBenchmarkRotatorSettings(b) wg sync.WaitGroup - errWg sync.WaitGroup err error ) @@ -1999,14 +1931,6 @@ func BenchmarkConcurrentOver2MB(b *testing.B) { b.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - errWg.Add(1) - go func() { - defer errWg.Done() - for err := range httpClient.ErrChan { - b.Errorf("Error writing to WARC: %s", err.Err.Error()) - } - }() - wg.Add(b.N) for n := 0; n < b.N; n++ { go func() { @@ -2014,13 +1938,11 @@ func BenchmarkConcurrentOver2MB(b *testing.B) { req, err := http.NewRequest("GET", server.URL, nil) if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } resp, err := httpClient.Do(req) if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } defer resp.Body.Close() @@ -2037,7 +1959,6 @@ func BenchmarkConcurrentOver2MBZStandard(b *testing.B) { var ( rotatorSettings = defaultBenchmarkRotatorSettings(b) wg sync.WaitGroup - errWg sync.WaitGroup err error ) rotatorSettings.Compression = "ZSTD" @@ -2061,14 +1982,6 @@ func BenchmarkConcurrentOver2MBZStandard(b *testing.B) { b.Fatalf("Unable to init WARC writing HTTP client: %s", err) } - errWg.Add(1) - go func() { - defer errWg.Done() - for err := range httpClient.ErrChan { - b.Errorf("Error writing to WARC: %s", err.Err.Error()) - } - }() - wg.Add(b.N) for n := 0; n < b.N; n++ { go func() { @@ -2076,13 +1989,11 @@ func BenchmarkConcurrentOver2MBZStandard(b *testing.B) { req, err := http.NewRequest("GET", server.URL, nil) if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } resp, err := httpClient.Do(req) if err != nil { - httpClient.ErrChan <- &Error{Err: err} return } defer resp.Body.Close() diff --git a/dedupe.go b/dedupe.go index ba1b36b..c93cb94 100644 --- a/dedupe.go +++ b/dedupe.go @@ -51,7 +51,7 @@ func (d *customDialer) checkLocalRevisit(digest string) revisitRecord { func checkCDXRevisit(CDXURL string, digest string, targetURI string, cookie string) (revisitRecord, error) { // CDX expects no hash header. For now we need to strip it. digest = strings.SplitN(digest, ":", 2)[1] - + req, err := http.NewRequest("GET", CDXURL+"/web/timemap/cdx?url="+url.QueryEscape(targetURI)+"&limit=-1", nil) if err != nil { return revisitRecord{}, err @@ -95,7 +95,7 @@ func checkCDXRevisit(CDXURL string, digest string, targetURI string, cookie stri func checkDoppelgangerRevisit(DoppelgangerHost string, digest string, targetURI string) (revisitRecord, error) { // Doppelganger is not expecting a hash header either but this will all be rewritten ... shortly... digest = strings.SplitN(digest, ":", 2)[1] - + req, err := http.NewRequest("GET", DoppelgangerHost+"/api/records/"+digest+"?uri="+targetURI, nil) if err != nil { return revisitRecord{}, err diff --git a/dialer.go b/dialer.go index c642cc5..8108330 100644 --- a/dialer.go +++ b/dialer.go @@ -9,6 +9,7 @@ import ( "net" "net/http" "net/url" + "path/filepath" "slices" "strconv" "strings" @@ -39,6 +40,12 @@ const ( // This is used internally to retrieve the wrapped connection for advanced use cases. // Use WithWrappedConnection() helper function for convenience. ContextKeyWrappedConn contextKey = "wrappedConn" + + // ContextKeyProxyType is the context key for requesting a specific proxy type. + // External callers (like Zeno) can set this to request a proxy of a specific type + // (Mobile, Residential, or Datacenter). + // Use WithProxyType() helper function for convenience. + ContextKeyProxyType contextKey = "proxyType" ) // WithFeedbackChannel adds a feedback channel to the request context. @@ -62,23 +69,51 @@ func WithWrappedConnection(ctx context.Context, wrappedConnChan chan *CustomConn return context.WithValue(ctx, ContextKeyWrappedConn, wrappedConnChan) } +// WithProxyType adds a proxy type preference to the request context. +// When set, the proxy selector will prefer proxies of the specified type +// (Mobile, Residential, or Datacenter). +// This is typically used by external callers like Zeno when they need a specific proxy type. +// +// Example: +// +// req = req.WithContext(warc.WithProxyType(req.Context(), warc.ProxyTypeMobile)) +func WithProxyType(ctx context.Context, proxyType ProxyType) context.Context { + return context.WithValue(ctx, ContextKeyProxyType, proxyType) +} + // dnsExchanger is an interface for DNS clients that can exchange messages type dnsExchanger interface { ExchangeContext(ctx context.Context, m *dns.Msg, address string) (r *dns.Msg, rtt time.Duration, err error) } +// proxyDialerInfo holds information about a configured proxy +type proxyDialerInfo struct { + dialer proxy.ContextDialer + needsHostname bool // true if proxy requires hostname (socks5h, http), false if can use IP (socks5) + proxyNetwork ProxyNetwork + proxyType ProxyType + allowedDomains []string // glob patterns + url string + name string + stats StatsRegistry +} + type customDialer struct { - proxyDialer proxy.ContextDialer - proxyNeedsHostname bool // true if proxy requires hostname (socks5h, http), false if can use IP (socks5) - client *CustomHTTPClient - DNSConfig *dns.ClientConfig - DNSClient dnsExchanger - DNSRecords *otter.Cache[string, net.IP] + proxyDialers []proxyDialerInfo + proxyRoundRobinIndex atomic.Uint32 + allowDirectFallback bool + client *CustomHTTPClient + DNSConfig *dns.ClientConfig + DNSClient dnsExchanger + DNSRecords *otter.Cache[string, net.IP] net.Dialer disableIPv4 bool disableIPv6 bool dnsConcurrency int dnsRoundRobinIndex atomic.Uint32 + + stats StatsRegistry + logBackend LogBackend } var emptyPayloadDigests = []string{ @@ -88,14 +123,18 @@ var emptyPayloadDigests = []string{ "blake3:af1349b9f5f9a1a6a0404dea36dcc9499bcb25c9adc112b7cc9a93cae41f3262", } -func newCustomDialer(httpClient *CustomHTTPClient, proxyURL string, DialTimeout, DNSRecordsTTL, DNSResolutionTimeout time.Duration, DNSCacheSize int, DNSServers []string, DNSConcurrency int, disableIPv4, disableIPv6 bool) (d *customDialer, err error) { +func newCustomDialer(httpClient *CustomHTTPClient, proxies []ProxyConfig, allowDirectFallback bool, DialTimeout, DNSRecordsTTL, DNSResolutionTimeout time.Duration, DNSCacheSize int, DNSServers []string, DNSConcurrency int, disableIPv4, disableIPv6 bool) (d *customDialer, err error) { d = new(customDialer) + d.stats = httpClient.statsRegistry + d.logBackend = httpClient.logBackend + d.Timeout = DialTimeout d.client = httpClient d.disableIPv4 = disableIPv4 d.disableIPv6 = disableIPv6 d.dnsConcurrency = DNSConcurrency + d.allowDirectFallback = allowDirectFallback DNScache, err := otter.MustBuilder[string, net.IP](DNSCacheSize). // CollectStats(). // Uncomment this line to enable stats collection, can be useful later on @@ -121,29 +160,163 @@ func newCustomDialer(httpClient *CustomHTTPClient, proxyURL string, DialTimeout, Timeout: DNSResolutionTimeout, } - if proxyURL != "" { - u, err := url.Parse(proxyURL) + // Initialize all proxies + for _, proxyConfig := range proxies { + if proxyConfig.URL == "" { + continue + } + + // Validate that Network is explicitly set + if proxyConfig.Network == ProxyNetworkUnset { + return nil, fmt.Errorf("proxy %s: Network must be explicitly set to ProxyNetworkAny, ProxyNetworkIPv4, or ProxyNetworkIPv6", proxyConfig.URL) + } + + u, err := url.Parse(proxyConfig.URL) if err != nil { - return nil, err + return nil, fmt.Errorf("failed to parse proxy URL %s: %w", proxyConfig.URL, err) } var proxyDialer proxy.Dialer if proxyDialer, err = proxy.FromURL(u, d); err != nil { - return nil, err + return nil, fmt.Errorf("failed to create proxy from URL %s: %w", proxyConfig.URL, err) } - d.proxyDialer = proxyDialer.(proxy.ContextDialer) - // Determine if this proxy requires hostname (remote DNS) or can use IP (local DNS) // Proxies with remote DNS: socks5h, socks4a, http, https // Proxies with local DNS: socks5, socks4 - d.proxyNeedsHostname = u.Scheme == "socks5h" || u.Scheme == "socks4a" || + needsHostname := u.Scheme == "socks5h" || u.Scheme == "socks4a" || u.Scheme == "http" || u.Scheme == "https" + + d.proxyDialers = append(d.proxyDialers, proxyDialerInfo{ + dialer: proxyDialer.(proxy.ContextDialer), + needsHostname: needsHostname, + proxyNetwork: proxyConfig.Network, + proxyType: proxyConfig.Type, + allowedDomains: proxyConfig.AllowedDomains, + url: proxyConfig.URL, + name: proxyName(u), + stats: httpClient.statsRegistry, + }) } return d, nil } +// selectProxy selects an appropriate proxy based on network type, domain, and context flags. +// Returns nil if no proxy is available or should be used (direct connection). +// Returns an error if proxies exist but none match the requirements and direct fallback is disabled. +func (d *customDialer) selectProxy(ctx context.Context, network, address string) (*proxyDialerInfo, error) { + // No proxies configured, use direct connection + if len(d.proxyDialers) == 0 { + return nil, nil + } + + // Extract hostname from address for domain matching + hostname, _, err := net.SplitHostPort(address) + if err != nil { + // If no port, treat the whole address as hostname + hostname = address + } + + // Check if a specific proxy type is requested via context + var requestedProxyType *ProxyType + if ctx.Value(ContextKeyProxyType) != nil { + if val, ok := ctx.Value(ContextKeyProxyType).(ProxyType); ok { + requestedProxyType = &val + } + } + + // Filter eligible proxies + var eligible []*proxyDialerInfo + for i := range d.proxyDialers { + proxy := &d.proxyDialers[i] + + // Filter by proxy type (Mobile, Residential, Datacenter) + // If a specific type is requested, only use matching proxies + // If no type is requested, only use ProxyTypeAny proxies + if requestedProxyType != nil { + // Specific type requested: only match that exact type + if proxy.proxyType != *requestedProxyType { + continue + } + } else { + // No type requested: only use ProxyTypeAny proxies + if proxy.proxyType != ProxyTypeAny { + continue + } + } + + // Filter by network type (IPv4/IPv6) + switch proxy.proxyNetwork { + case ProxyNetworkIPv4: + if strings.HasSuffix(network, "6") { + continue // Skip IPv6 networks + } + case ProxyNetworkIPv6: + if strings.HasSuffix(network, "4") { + continue // Skip IPv4 networks + } + case ProxyNetworkAny: + // Always eligible regardless of network type + } + + // Filter by domain patterns + if len(proxy.allowedDomains) > 0 { + matched := false + for _, pattern := range proxy.allowedDomains { + // Use filepath.Match for glob pattern matching + // Note: filepath.Match doesn't support '**' but supports '*' and '?' + if match, _ := filepath.Match(pattern, hostname); match { + matched = true + break + } + // Also check if pattern matches a subdomain pattern + // For example, "*.example.com" should match "api.example.com" + if strings.HasPrefix(pattern, "*.") { + suffix := pattern[1:] // Remove the leading '*' + if strings.HasSuffix(hostname, suffix) || hostname == suffix[1:] { + matched = true + break + } + } + } + if !matched { + continue // Skip if domain doesn't match any pattern + } + } + + eligible = append(eligible, proxy) + } + + // No eligible proxies found + if len(eligible) == 0 { + if d.allowDirectFallback { + d.logBackend.Debug("no eligible proxies found, using direct connection", "network", network, "address", address) + return nil, nil // Use direct connection + } + proxyTypeStr := "any" + if requestedProxyType != nil { + proxyTypeStr = fmt.Sprintf("%v", *requestedProxyType) + } + d.logBackend.Error("no eligible proxies found and direct fallback disabled", "network", network, "address", address, "proxyType", proxyTypeStr) + return nil, fmt.Errorf("no eligible proxies found for network=%s, address=%s, proxyType=%s and direct fallback is disabled", network, address, proxyTypeStr) + } + + // Round-robin selection among eligible proxies + startIdx := int(d.proxyRoundRobinIndex.Add(1)-1) % len(eligible) + selectedProxy := eligible[startIdx] + + // Update proxy statistics + if selectedProxy.stats != nil { + selectedProxy.stats.RegisterCounter(proxyRequestsTotal, proxyRequestsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": selectedProxy.name}).Add(1) + selectedProxy.stats.RegisterGauge(proxyLastUsedNanoseconds, proxyLastUsedNanosecondsHelp, []string{"proxy"}).WithLabels(Labels{"proxy": selectedProxy.name}).Set(time.Now().UnixNano()) + } + + d.logBackend.Debug("proxy selected", "proxy", selectedProxy.name, "network", network, "address", address) + + return selectedProxy, nil +} + type CustomConnection struct { net.Conn io.Reader @@ -235,10 +408,16 @@ func (d *customDialer) CustomDialContext(ctx context.Context, network, address s return nil, errors.New("no supported network type available") } + // Select appropriate proxy based on context, network type, and domain + selectedProxy, err := d.selectProxy(ctx, network, address) + if err != nil { + return nil, err + } + var dialAddr string var IP net.IP - if d.proxyDialer != nil && d.proxyNeedsHostname { + if selectedProxy != nil && selectedProxy.needsHostname { // Remote DNS proxy (socks5h, socks4a, http, https) // Skip DNS archiving to avoid privacy leak and ensure accuracy. // The proxy will handle DNS resolution on its end, and we don't want to: @@ -250,6 +429,7 @@ func (d *customDialer) CustomDialContext(ctx context.Context, network, address s // Archive DNS and use resolved IP IP, _, err = d.archiveDNS(ctx, address) if err != nil { + d.logBackend.Error("DNS resolution failed", "address", address, "error", err) return nil, err } @@ -261,10 +441,19 @@ func (d *customDialer) CustomDialContext(ctx context.Context, network, address s } dialAddr = net.JoinHostPort(IP.String(), port) + d.logBackend.Debug("DNS resolved", "address", address, "ip", IP.String()) } - if d.proxyDialer != nil { - conn, err = d.proxyDialer.DialContext(ctx, network, dialAddr) + if selectedProxy != nil { + conn, err = selectedProxy.dialer.DialContext(ctx, network, dialAddr) + if err != nil { + d.logBackend.Error("proxy connection failed", "proxy", selectedProxy.name, "address", dialAddr, "error", err) + if selectedProxy.stats != nil { + selectedProxy.stats.RegisterCounter(proxyErrorsTotal, proxyErrorsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": selectedProxy.name}).Add(1) + } + } else { + d.logBackend.Debug("connection established via proxy", "proxy", selectedProxy.name, "address", dialAddr) + } } else { if d.client.randomLocalIP { localAddr := getLocalAddr(network, IP) @@ -279,6 +468,11 @@ func (d *customDialer) CustomDialContext(ctx context.Context, network, address s } conn, err = d.DialContext(ctx, network, dialAddr) + if err != nil { + d.logBackend.Error("direct connection failed", "address", dialAddr, "error", err) + } else { + d.logBackend.Debug("direct connection established", "address", dialAddr) + } } if err != nil { @@ -299,11 +493,16 @@ func (d *customDialer) CustomDialTLSContext(ctx context.Context, network, addres return nil, errors.New("no supported network type available") } + // Select appropriate proxy based on context, network type, and domain + selectedProxy, err := d.selectProxy(ctx, network, address) + if err != nil { + return nil, err + } + var dialAddr string var IP net.IP - var err error - if d.proxyDialer != nil && d.proxyNeedsHostname { + if selectedProxy != nil && selectedProxy.needsHostname { // Remote DNS proxy (socks5h, socks4a, http, https) // Skip DNS archiving to avoid privacy leak and ensure accuracy. // The proxy will handle DNS resolution on its end, and we don't want to: @@ -315,6 +514,7 @@ func (d *customDialer) CustomDialTLSContext(ctx context.Context, network, addres // Archive DNS and use resolved IP IP, _, err = d.archiveDNS(ctx, address) if err != nil { + d.logBackend.Error("DNS resolution failed for TLS connection", "address", address, "error", err) return nil, err } @@ -326,12 +526,16 @@ func (d *customDialer) CustomDialTLSContext(ctx context.Context, network, addres } dialAddr = net.JoinHostPort(IP.String(), port) + d.logBackend.Debug("DNS resolved for TLS connection", "address", address, "ip", IP.String()) } var plainConn net.Conn - if d.proxyDialer != nil { - plainConn, err = d.proxyDialer.DialContext(ctx, network, dialAddr) + if selectedProxy != nil { + plainConn, err = selectedProxy.dialer.DialContext(ctx, network, dialAddr) + if err != nil && selectedProxy.stats != nil { + selectedProxy.stats.RegisterCounter(proxyErrorsTotal, proxyErrorsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": selectedProxy.name}).Add(1) + } } else { if d.client.randomLocalIP { localAddr := getLocalAddr(network, IP) @@ -367,6 +571,15 @@ func (d *customDialer) CustomDialTLSContext(ctx context.Context, network, addres defer cancel() if err := tlsConn.HandshakeContext(handshakeCtx); err != nil { + // Track TLS handshake errors for proxy connections + if selectedProxy != nil { + d.logBackend.Error("TLS handshake failed via proxy", "proxy", selectedProxy.name, "address", address, "error", err) + if selectedProxy.stats != nil { + selectedProxy.stats.RegisterCounter(proxyErrorsTotal, proxyErrorsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": selectedProxy.name}).Add(1) + } + } else { + d.logBackend.Error("TLS handshake failed", "address", address, "error", err) + } closeErr := plainConn.Close() if closeErr != nil { return nil, fmt.Errorf("CustomDialTLS: TLS handshake failed and closing plain connection failed: %s", closeErr.Error()) @@ -374,6 +587,12 @@ func (d *customDialer) CustomDialTLSContext(ctx context.Context, network, addres return nil, fmt.Errorf("CustomDialTLS: TLS handshake failed: %w", err) } + if selectedProxy != nil { + d.logBackend.Debug("TLS connection established via proxy", "proxy", selectedProxy.name, "address", address) + } else { + d.logBackend.Debug("TLS connection established", "address", address) + } + return d.wrapConnection(ctx, tlsConn, "https"), nil } @@ -448,17 +667,11 @@ func (d *customDialer) writeWARCFromConnection(ctx context.Context, reqPipe, res close(recordChan) if readErr != nil { - d.client.ErrChan <- &Error{ - Err: readErr, - Func: "writeWARCFromConnection", - } + d.logBackend.Error("error reading from connection", "func", "writeWARCFromConnection", "error", readErr) for record := range recordChan { if closeErr := record.Content.Close(); closeErr != nil { - d.client.ErrChan <- &Error{ - Err: closeErr, - Func: "writeWARCFromConnection", - } + d.logBackend.Error("error closing record content", "func", "writeWARCFromConnection", "error", closeErr) } } @@ -477,14 +690,11 @@ func (d *customDialer) writeWARCFromConnection(ctx context.Context, reqPipe, res if len(batch.Records) != 2 { err.Err = errors.New("warc: there was an unspecified problem creating one of the WARC records") - d.client.ErrChan <- err + d.logBackend.Error("failed to create WARC records", "func", err.Func, "error", err.Err) for _, record := range batch.Records { if closeErr := record.Content.Close(); closeErr != nil { - d.client.ErrChan <- &Error{ - Err: closeErr, - Func: "writeWARCFromConnection", - } + d.logBackend.Error("error closing record content", "func", "writeWARCFromConnection", "error", closeErr) } } @@ -511,7 +721,7 @@ func (d *customDialer) writeWARCFromConnection(ctx context.Context, reqPipe, res case <-ctx.Done(): return default: - if d.proxyDialer == nil { + if len(d.proxyDialers) == 0 { switch addr := conn.RemoteAddr().(type) { case *net.TCPAddr: IP := addr.IP.String() @@ -530,19 +740,13 @@ func (d *customDialer) writeWARCFromConnection(ctx context.Context, reqPipe, res r.Header.Set("WARC-Target-URI", warcTargetURI) if _, seekErr := r.Content.Seek(0, 0); seekErr != nil { - d.client.ErrChan <- &Error{ - Err: seekErr, - Func: "writeWARCFromConnection", - } + d.logBackend.Error("error seeking record content", "func", "writeWARCFromConnection", "error", seekErr) return } digest, err := GetDigest(r.Content, d.client.DigestAlgorithm) if err != nil { - d.client.ErrChan <- &Error{ - Err: err, - Func: "writeWARCFromConnection", - } + d.logBackend.Error("error calculating digest", "func", "writeWARCFromConnection", "error", err) return } @@ -553,10 +757,7 @@ func (d *customDialer) writeWARCFromConnection(ctx context.Context, reqPipe, res if r.Header.Get("WARC-Type") == "response" && !slices.Contains(emptyPayloadDigests, r.Header.Get("WARC-Payload-Digest")) { captureTime, timeConversionErr := time.Parse(time.RFC3339, batch.CaptureTime) if timeConversionErr != nil { - d.client.ErrChan <- &Error{ - Err: timeConversionErr, - Func: "writeWARCFromConnection.timeConversionErr", - } + d.logBackend.Error("error parsing capture time", "func", "writeWARCFromConnection", "error", timeConversionErr) return } d.client.dedupeHashTable.Store(r.Header.Get("WARC-Payload-Digest"), revisitRecord{ @@ -645,8 +846,8 @@ func (d *customDialer) readResponse(ctx context.Context, respPipe *io.PipeReader if d.client.dedupeOptions.LocalDedupe { revisit = d.checkLocalRevisit(payloadDigest) if revisit.targetURI != "" { - LocalDedupeTotalBytes.Add(int64(revisit.size)) - LocalDedupeTotal.Add(1) + d.stats.RegisterCounter(localDedupedBytesTotal, localDedupedBytesTotalHelp, nil).WithLabels(nil).Add(int64(revisit.size)) + d.stats.RegisterCounter(localDedupedTotal, localDedupedTotalHelp, nil).WithLabels(nil).Add(1) } } @@ -655,8 +856,8 @@ func (d *customDialer) readResponse(ctx context.Context, respPipe *io.PipeReader if d.client.dedupeOptions.DoppelgangerDedupe && d.client.DigestAlgorithm == SHA1 && revisit.targetURI == "" { revisit, _ = checkDoppelgangerRevisit(d.client.dedupeOptions.DoppelgangerHost, payloadDigest, warcTargetURI) if revisit.targetURI != "" { - DoppelgangerDedupeTotalBytes.Add(bytesCopied) - DoppelgangerDedupeTotal.Add(1) + d.stats.RegisterCounter(doppelgangerDedupedBytesTotal, doppelgangerDedupedBytesTotalHelp, nil).WithLabels(nil).Add(bytesCopied) + d.stats.RegisterCounter(doppelgangerDedupedTotal, doppelgangerDedupedTotalHelp, nil).WithLabels(nil).Add(1) } } @@ -664,8 +865,8 @@ func (d *customDialer) readResponse(ctx context.Context, respPipe *io.PipeReader if d.client.dedupeOptions.CDXDedupe && d.client.DigestAlgorithm == SHA1 && revisit.targetURI == "" { revisit, _ = checkCDXRevisit(d.client.dedupeOptions.CDXURL, payloadDigest, warcTargetURI, d.client.dedupeOptions.CDXCookie) if revisit.targetURI != "" { - CDXDedupeTotalBytes.Add(bytesCopied) - CDXDedupeTotal.Add(1) + d.stats.RegisterCounter(cdxDedupedBytesTotal, cdxDedupedBytesTotalHelp, nil).WithLabels(nil).Add(bytesCopied) + d.stats.RegisterCounter(cdxDedupedTotal, cdxDedupedTotalHelp, nil).WithLabels(nil).Add(1) } } } diff --git a/dialer_test.go b/dialer_test.go index 147668b..0baff06 100644 --- a/dialer_test.go +++ b/dialer_test.go @@ -2,9 +2,11 @@ package warc import ( "bytes" + "context" "io" "strings" "testing" + "time" ) func TestGetNetworkType(t *testing.T) { @@ -165,3 +167,713 @@ func TestFindEndOfHeadersOffset(t *testing.T) { }) } } + +func TestProxySelection(t *testing.T) { + t.Run("NoProxies", func(t *testing.T) { + d := &customDialer{ + proxyDialers: []proxyDialerInfo{}, + logBackend: &noopLogger{}, + } + proxy, err := d.selectProxy(context.Background(), "tcp", "example.com:80") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy != nil { + t.Error("expected nil proxy when no proxies configured") + } + }) + + t.Run("IPv4ProxyWithIPv4Network", func(t *testing.T) { + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkIPv4, + proxyType: ProxyTypeAny, + url: "socks5://ipv4-proxy:1080", + }, + }, + logBackend: &noopLogger{}, + } + proxy, err := d.selectProxy(context.Background(), "tcp4", "example.com:80") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy == nil { + t.Error("expected proxy for IPv4 network with IPv4 proxy") + } + if proxy != nil && proxy.url != "socks5://ipv4-proxy:1080" { + t.Errorf("expected socks5://ipv4-proxy:1080, got %s", proxy.url) + } + }) + + t.Run("IPv4ProxyWithIPv6Network", func(t *testing.T) { + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkIPv4, + proxyType: ProxyTypeAny, + url: "socks5://ipv4-proxy:1080", + }, + }, + logBackend: &noopLogger{}, + allowDirectFallback: true, + } + proxy, err := d.selectProxy(context.Background(), "tcp6", "example.com:80") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy != nil { + t.Error("expected nil proxy for IPv6 network with IPv4 proxy") + } + }) + + t.Run("IPv6ProxyWithIPv6Network", func(t *testing.T) { + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkIPv6, + proxyType: ProxyTypeAny, + url: "socks5://ipv6-proxy:1080", + }, + }, + logBackend: &noopLogger{}, + } + proxy, err := d.selectProxy(context.Background(), "tcp6", "example.com:80") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy == nil { + t.Error("expected proxy for IPv6 network with IPv6 proxy") + } + }) + + t.Run("DomainFiltering", func(t *testing.T) { + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + allowedDomains: []string{"*.example.com"}, + url: "socks5://domain-proxy:1080", + }, + }, + logBackend: &noopLogger{}, + } + + // Should match subdomain + proxy, err := d.selectProxy(context.Background(), "tcp", "api.example.com:443") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy == nil { + t.Error("expected proxy for matching domain") + } + + // Should match base domain + d.allowDirectFallback = true + proxy, err = d.selectProxy(context.Background(), "tcp", "example.com:443") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy == nil { + t.Error("expected proxy for base domain") + } + + // Should not match different domain + proxy, err = d.selectProxy(context.Background(), "tcp", "other.com:443") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy != nil { + t.Error("expected nil proxy for non-matching domain") + } + }) + + t.Run("ProxyTypeSelection", func(t *testing.T) { + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + url: "socks5://any-proxy:1080", + }, + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeMobile, + url: "socks5://mobile-proxy:1080", + }, + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeResidential, + url: "socks5://residential-proxy:1080", + }, + }, + logBackend: &noopLogger{}, + } + + // Without proxy type context, should use ProxyTypeAny proxy + proxy, err := d.selectProxy(context.Background(), "tcp", "example.com:80") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy == nil || proxy.url != "socks5://any-proxy:1080" { + t.Error("expected any-proxy without proxy type context") + } + + // With mobile proxy type context, should use mobile proxy + ctx := WithProxyType(context.Background(), ProxyTypeMobile) + proxy, err = d.selectProxy(ctx, "tcp", "example.com:80") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy == nil || proxy.url != "socks5://mobile-proxy:1080" { + t.Error("expected mobile-proxy with mobile proxy type context") + } + + // With residential proxy type context, should use residential proxy + ctx = WithProxyType(context.Background(), ProxyTypeResidential) + proxy, err = d.selectProxy(ctx, "tcp", "example.com:80") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy == nil || proxy.url != "socks5://residential-proxy:1080" { + t.Error("expected residential-proxy with residential proxy type context") + } + }) + + t.Run("RoundRobinSelection", func(t *testing.T) { + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + url: "socks5://proxy1:1080", + }, + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + url: "socks5://proxy2:1080", + }, + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + url: "socks5://proxy3:1080", + }, + }, + logBackend: &noopLogger{}, + } + + // Expected order for 3 complete cycles (9 selections) + expectedOrder := []string{ + "socks5://proxy1:1080", + "socks5://proxy2:1080", + "socks5://proxy3:1080", + "socks5://proxy1:1080", + "socks5://proxy2:1080", + "socks5://proxy3:1080", + "socks5://proxy1:1080", + "socks5://proxy2:1080", + "socks5://proxy3:1080", + } + + // Select proxies 9 times and verify sequential round-robin order + for i := 0; i < 9; i++ { + proxy, err := d.selectProxy(context.Background(), "tcp", "example.com:80") + if err != nil { + t.Errorf("iteration %d: unexpected error: %v", i, err) + } + if proxy == nil { + t.Errorf("iteration %d: expected proxy, got nil", i) + } else if proxy.url != expectedOrder[i] { + t.Errorf("iteration %d: expected %s, got %s", i, expectedOrder[i], proxy.url) + } + } + }) + + t.Run("NoEligibleProxiesWithFallback", func(t *testing.T) { + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkIPv6, + proxyType: ProxyTypeAny, + url: "socks5://ipv6-proxy:1080", + }, + }, + logBackend: &noopLogger{}, + allowDirectFallback: true, + } + + // IPv4 network with IPv6 proxy should use direct connection + proxy, err := d.selectProxy(context.Background(), "tcp4", "example.com:80") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy != nil { + t.Error("expected nil proxy with direct fallback") + } + }) + + t.Run("NoEligibleProxiesWithoutFallback", func(t *testing.T) { + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkIPv6, + proxyType: ProxyTypeAny, + url: "socks5://ipv6-proxy:1080", + }, + }, + logBackend: &noopLogger{}, + allowDirectFallback: false, + } + + // IPv4 network with IPv6 proxy and no fallback should error + proxy, err := d.selectProxy(context.Background(), "tcp4", "example.com:80") + if err == nil { + t.Error("expected error when no eligible proxies and no fallback") + } + if proxy != nil { + t.Error("expected nil proxy") + } + }) + + t.Run("ComplexFiltering", func(t *testing.T) { + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkIPv4, + proxyType: ProxyTypeAny, + allowedDomains: []string{"*.api.example.com"}, + url: "socks5://api-ipv4-proxy:1080", + }, + { + proxyNetwork: ProxyNetworkIPv6, + proxyType: ProxyTypeAny, + allowedDomains: []string{"*.media.example.com"}, + url: "socks5://media-ipv6-proxy:1080", + }, + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeMobile, + url: "socks5://mobile-proxy:1080", + }, + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeResidential, + url: "socks5://residential-proxy:1080", + }, + }, + logBackend: &noopLogger{}, + } + + // Test IPv4 API domain + proxy, err := d.selectProxy(context.Background(), "tcp4", "service.api.example.com:443") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy == nil || proxy.url != "socks5://api-ipv4-proxy:1080" { + t.Error("expected api-ipv4-proxy for IPv4 API domain") + } + + // Test IPv6 media domain + proxy, err = d.selectProxy(context.Background(), "tcp6", "cdn.media.example.com:443") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy == nil { + t.Error("expected proxy for IPv6 media domain, got nil") + } else if proxy.url != "socks5://media-ipv6-proxy:1080" { + t.Errorf("expected media-ipv6-proxy for IPv6 media domain, got %s", proxy.url) + } + + // Test mobile proxy type context + ctx := WithProxyType(context.Background(), ProxyTypeMobile) + proxy, err = d.selectProxy(ctx, "tcp", "example.com:80") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy == nil || proxy.url != "socks5://mobile-proxy:1080" { + t.Error("expected mobile-proxy with mobile proxy type context") + } + + // Test residential proxy type context + ctx = WithProxyType(context.Background(), ProxyTypeResidential) + proxy, err = d.selectProxy(ctx, "tcp", "example.com:80") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if proxy == nil || proxy.url != "socks5://residential-proxy:1080" { + t.Error("expected residential-proxy with residential proxy type context") + } + }) +} + +// TestProxyStatsMetricNames tests that proxy metrics use labels to distinguish between proxies +func TestProxyStatsMetricNames(t *testing.T) { + registry := newLocalRegistry() + + tests := []struct { + name string + proxyName string + }{ + { + name: "simple proxy", + proxyName: "example_com_8080", + }, + { + name: "IPv4 proxy", + proxyName: "192_168_1_1_3128", + }, + { + name: "IPv6 proxy", + proxyName: "2001_db8__1_8080", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Register counters for this proxy using labels + requestsCounter := registry.RegisterCounter(proxyRequestsTotal, proxyRequestsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": tt.proxyName}) + errorsCounter := registry.RegisterCounter(proxyErrorsTotal, proxyErrorsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": tt.proxyName}) + lastUsedGauge := registry.RegisterGauge(proxyLastUsedNanoseconds, proxyLastUsedNanosecondsHelp, []string{"proxy"}).WithLabels(Labels{"proxy": tt.proxyName}) + + // Verify counters start at 0 + if requestsCounter.Get() != 0 { + t.Errorf("Expected requests counter to start at 0, got %d", requestsCounter.Get()) + } + if errorsCounter.Get() != 0 { + t.Errorf("Expected errors counter to start at 0, got %d", errorsCounter.Get()) + } + + // Increment counters + requestsCounter.Add(5) + errorsCounter.Add(2) + lastUsedGauge.Set(123456789) + + // Verify values are independent per proxy + if requestsCounter.Get() != 5 { + t.Errorf("Expected requests counter to be 5, got %d", requestsCounter.Get()) + } + if errorsCounter.Get() != 2 { + t.Errorf("Expected errors counter to be 2, got %d", errorsCounter.Get()) + } + if lastUsedGauge.Get() != 123456789 { + t.Errorf("Expected last used gauge to be 123456789, got %d", lastUsedGauge.Get()) + } + }) + } + + // Verify that different proxy labels create independent counters + proxy1Counter := registry.RegisterCounter(proxyRequestsTotal, proxyRequestsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": "example_com_8080"}) + proxy2Counter := registry.RegisterCounter(proxyRequestsTotal, proxyRequestsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": "192_168_1_1_3128"}) + + if proxy1Counter.Get() != 5 { + t.Errorf("Expected proxy1 counter to be 5 (from earlier test), got %d", proxy1Counter.Get()) + } + if proxy2Counter.Get() != 5 { + t.Errorf("Expected proxy2 counter to be 5 (from earlier test), got %d", proxy2Counter.Get()) + } +} + +// TestProxyStatsRequestCount tests that proxy request counts are incremented correctly +func TestProxyStatsRequestCount(t *testing.T) { + registry := newLocalRegistry() + + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + url: "socks5://proxy1:1080", + name: "proxy1_1080", + stats: registry, + }, + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + url: "socks5://proxy2:1080", + name: "proxy2_1080", + stats: registry, + }, + }, + logBackend: &noopLogger{}, + } + + // Select proxies multiple times and verify request counts + for i := 0; i < 5; i++ { + proxy, err := d.selectProxy(context.Background(), "tcp", "example.com:80") + if err != nil { + t.Fatalf("iteration %d: unexpected error: %v", i, err) + } + if proxy == nil { + t.Fatalf("iteration %d: expected proxy, got nil", i) + } + } + + // Verify both proxies were used (round-robin) + // With 5 selections: proxy1 should be used 3 times, proxy2 should be used 2 times + proxy1Counter := registry.RegisterCounter(proxyRequestsTotal, proxyRequestsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": "proxy1_1080"}) + proxy2Counter := registry.RegisterCounter(proxyRequestsTotal, proxyRequestsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": "proxy2_1080"}) + + if proxy1Counter.Get() != 3 { + t.Errorf("Expected proxy1 request count 3, got %d", proxy1Counter.Get()) + } + if proxy2Counter.Get() != 2 { + t.Errorf("Expected proxy2 request count 2, got %d", proxy2Counter.Get()) + } +} + +// TestProxyStatsLastUsed tests that proxy last used timestamps are updated +func TestProxyStatsLastUsed(t *testing.T) { + registry := newLocalRegistry() + + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + url: "socks5://proxy:1080", + name: "proxy_1080", + stats: registry, + }, + }, + logBackend: &noopLogger{}, + } + + // Record time before selection + timeBefore := time.Now().UnixNano() + + // Select proxy + proxy, err := d.selectProxy(context.Background(), "tcp", "example.com:80") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if proxy == nil { + t.Fatal("expected proxy, got nil") + } + + // Record time after selection + timeAfter := time.Now().UnixNano() + + // Verify last used timestamp is within expected range + lastUsedGauge := registry.RegisterGauge(proxyLastUsedNanoseconds, proxyLastUsedNanosecondsHelp, []string{"proxy"}).WithLabels(Labels{"proxy": "proxy_1080"}) + lastUsed := lastUsedGauge.Get() + + if lastUsed < timeBefore || lastUsed > timeAfter { + t.Errorf("Expected last used timestamp between %d and %d, got %d", timeBefore, timeAfter, lastUsed) + } + + // Wait a bit and select again + time.Sleep(10 * time.Millisecond) + timeBeforeSecond := time.Now().UnixNano() + + proxy, err = d.selectProxy(context.Background(), "tcp", "example.com:80") + if err != nil { + t.Fatalf("unexpected error on second selection: %v", err) + } + if proxy == nil { + t.Fatal("expected proxy on second selection, got nil") + } + + timeAfterSecond := time.Now().UnixNano() + + // Verify last used timestamp was updated + lastUsedSecond := lastUsedGauge.Get() + if lastUsedSecond < timeBeforeSecond || lastUsedSecond > timeAfterSecond { + t.Errorf("Expected updated last used timestamp between %d and %d, got %d", timeBeforeSecond, timeAfterSecond, lastUsedSecond) + } + if lastUsedSecond <= lastUsed { + t.Errorf("Expected last used timestamp to be updated, but %d <= %d", lastUsedSecond, lastUsed) + } +} + +// TestProxyStatsWithNilRegistry tests that proxy selection works when stats registry is nil +func TestProxyStatsWithNilRegistry(t *testing.T) { + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + url: "socks5://proxy:1080", + name: "proxy_1080", + stats: nil, // No stats registry + }, + }, + logBackend: &noopLogger{}, + } + + // Should not panic when stats is nil + proxy, err := d.selectProxy(context.Background(), "tcp", "example.com:80") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if proxy == nil { + t.Fatal("expected proxy, got nil") + } +} + +// TestProxyStatsMultipleProxiesRoundRobin tests that stats are correctly tracked across multiple proxies in round-robin +func TestProxyStatsMultipleProxiesRoundRobin(t *testing.T) { + registry := newLocalRegistry() + + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + url: "socks5://proxy1:1080", + name: "proxy1", + stats: registry, + }, + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + url: "socks5://proxy2:1080", + name: "proxy2", + stats: registry, + }, + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + url: "socks5://proxy3:1080", + name: "proxy3", + stats: registry, + }, + }, + logBackend: &noopLogger{}, + } + + // Select proxies 12 times (4 complete round-robin cycles) + for i := 0; i < 12; i++ { + proxy, err := d.selectProxy(context.Background(), "tcp", "example.com:80") + if err != nil { + t.Fatalf("iteration %d: unexpected error: %v", i, err) + } + if proxy == nil { + t.Fatalf("iteration %d: expected proxy, got nil", i) + } + } + + // Verify each proxy was used exactly 4 times + for i := 1; i <= 3; i++ { + proxyName := "proxy" + string(rune('0'+i)) + counter := registry.RegisterCounter(proxyRequestsTotal, proxyRequestsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": proxyName}) + + expectedCount := int64(4) + if counter.Get() != expectedCount { + t.Errorf("Expected %s request count %d, got %d", proxyName, expectedCount, counter.Get()) + } + } +} + +// TestProxyStatsWithDomainFiltering tests that stats are only updated for eligible proxies +func TestProxyStatsWithDomainFiltering(t *testing.T) { + registry := newLocalRegistry() + + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + allowedDomains: []string{"*.example.com"}, + url: "socks5://example-proxy:1080", + name: "example_proxy", + stats: registry, + }, + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeAny, + allowedDomains: []string{"*.test.com"}, + url: "socks5://test-proxy:1080", + name: "test_proxy", + stats: registry, + }, + }, + logBackend: &noopLogger{}, + allowDirectFallback: true, + } + + // Select proxy for example.com domain - should use example-proxy + proxy, err := d.selectProxy(context.Background(), "tcp", "api.example.com:443") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if proxy == nil || proxy.name != "example_proxy" { + t.Fatal("expected example-proxy") + } + + // Select proxy for test.com domain - should use test-proxy + proxy, err = d.selectProxy(context.Background(), "tcp", "api.test.com:443") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if proxy == nil || proxy.name != "test_proxy" { + t.Fatal("expected test-proxy") + } + + // Verify stats + exampleCounter := registry.RegisterCounter(proxyRequestsTotal, proxyRequestsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": "example_proxy"}) + testCounter := registry.RegisterCounter(proxyRequestsTotal, proxyRequestsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": "test_proxy"}) + + if exampleCounter.Get() != 1 { + t.Errorf("Expected example_proxy request count 1, got %d", exampleCounter.Get()) + } + if testCounter.Get() != 1 { + t.Errorf("Expected test_proxy request count 1, got %d", testCounter.Get()) + } +} + +// TestProxyStatsProxyTypeFiltering tests that stats work correctly with proxy type filtering +func TestProxyStatsProxyTypeFiltering(t *testing.T) { + registry := newLocalRegistry() + + d := &customDialer{ + proxyDialers: []proxyDialerInfo{ + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeMobile, + url: "socks5://mobile-proxy:1080", + name: "mobile_proxy", + stats: registry, + }, + { + proxyNetwork: ProxyNetworkAny, + proxyType: ProxyTypeResidential, + url: "socks5://residential-proxy:1080", + name: "residential_proxy", + stats: registry, + }, + }, + logBackend: &noopLogger{}, + } + + // Select mobile proxy + ctx := WithProxyType(context.Background(), ProxyTypeMobile) + proxy, err := d.selectProxy(ctx, "tcp", "example.com:80") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if proxy == nil || proxy.proxyType != ProxyTypeMobile { + t.Fatal("expected mobile proxy") + } + + // Select residential proxy + ctx = WithProxyType(context.Background(), ProxyTypeResidential) + proxy, err = d.selectProxy(ctx, "tcp", "example.com:80") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if proxy == nil || proxy.proxyType != ProxyTypeResidential { + t.Fatal("expected residential proxy") + } + + // Verify stats + mobileCounter := registry.RegisterCounter(proxyRequestsTotal, proxyRequestsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": "mobile_proxy"}) + residentialCounter := registry.RegisterCounter(proxyRequestsTotal, proxyRequestsTotalHelp, []string{"proxy"}).WithLabels(Labels{"proxy": "residential_proxy"}) + + if mobileCounter.Get() != 1 { + t.Errorf("Expected mobile_proxy request count 1, got %d", mobileCounter.Get()) + } + if residentialCounter.Get() != 1 { + t.Errorf("Expected residential_proxy request count 1, got %d", residentialCounter.Get()) + } +} diff --git a/gzip_interface.go b/gzip_interface.go index 4ca2fa7..84ceb0f 100644 --- a/gzip_interface.go +++ b/gzip_interface.go @@ -19,4 +19,4 @@ type GzipReaderInterface interface { io.ReadCloser Multistream(enable bool) Reset(r io.Reader) error -} \ No newline at end of file +} diff --git a/logging.go b/logging.go new file mode 100644 index 0000000..376df6a --- /dev/null +++ b/logging.go @@ -0,0 +1,154 @@ +package warc + +import ( + "context" + "log/slog" + "sync" +) + +// LogBackend provides a pluggable logging interface compatible with slog. +// Users can implement this interface to integrate their preferred logging solution. +// The interface matches slog.Logger method signatures for easy wrapping. +// +// Users control the log level cutoff by configuring their logger implementation. +// For example, when wrapping slog.Logger, set the level in the handler: +// +// handler := slog.NewJSONHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelInfo}) +// logger := slog.New(handler) +// +// If no LogBackend is provided, a no-op logger is used by default. +type LogBackend interface { + // Debug logs a message at Debug level with optional key-value pairs + Debug(msg string, args ...any) + + // Info logs a message at Info level with optional key-value pairs + Info(msg string, args ...any) + + // Warn logs a message at Warn level with optional key-value pairs + Warn(msg string, args ...any) + + // Error logs a message at Error level with optional key-value pairs + Error(msg string, args ...any) +} + +// noopLogger is a no-op implementation of LogBackend used when no logger is provided. +type noopLogger struct{} + +func (n *noopLogger) Debug(_ string, _ ...any) {} +func (n *noopLogger) Info(_ string, _ ...any) {} +func (n *noopLogger) Warn(_ string, _ ...any) {} +func (n *noopLogger) Error(_ string, _ ...any) {} +func (n *noopLogger) Log(_ context.Context, _ slog.Level, _ string, _ ...any) {} + +// LogEntry represents a single log entry captured by TestLogger +type LogEntry struct { + Level slog.Level + Message string + Args []any +} + +// testLogger is a thread-safe logger implementation for unit tests that captures +// all log messages for verification. Messages can be retrieved sequentially. +type testLogger struct { + entries []LogEntry + mu sync.Mutex +} + +// NewTestLogger creates a new TestLogger for use in unit tests +func NewTestLogger() *testLogger { + return &testLogger{ + entries: make([]LogEntry, 0), + } +} + +func (t *testLogger) Debug(msg string, args ...any) { + t.mu.Lock() + defer t.mu.Unlock() + t.entries = append(t.entries, LogEntry{Level: slog.LevelDebug, Message: msg, Args: args}) +} + +func (t *testLogger) Info(msg string, args ...any) { + t.mu.Lock() + defer t.mu.Unlock() + t.entries = append(t.entries, LogEntry{Level: slog.LevelInfo, Message: msg, Args: args}) +} + +func (t *testLogger) Warn(msg string, args ...any) { + t.mu.Lock() + defer t.mu.Unlock() + t.entries = append(t.entries, LogEntry{Level: slog.LevelWarn, Message: msg, Args: args}) +} + +func (t *testLogger) Error(msg string, args ...any) { + t.mu.Lock() + defer t.mu.Unlock() + t.entries = append(t.entries, LogEntry{Level: slog.LevelError, Message: msg, Args: args}) +} + +func (t *testLogger) Log(ctx context.Context, level slog.Level, msg string, args ...any) { + t.mu.Lock() + defer t.mu.Unlock() + t.entries = append(t.entries, LogEntry{Level: level, Message: msg, Args: args}) +} + +// Entries returns all captured log entries +func (t *testLogger) Entries() []LogEntry { + t.mu.Lock() + defer t.mu.Unlock() + // Return a copy to prevent external modification + result := make([]LogEntry, len(t.entries)) + copy(result, t.entries) + return result +} + +// Next returns the next log entry and removes it from the queue. +// Returns nil if no entries are available. +func (t *testLogger) Next() *LogEntry { + t.mu.Lock() + defer t.mu.Unlock() + if len(t.entries) == 0 { + return nil + } + entry := t.entries[0] + t.entries = t.entries[1:] + return &entry +} + +// Count returns the number of captured log entries +func (t *testLogger) Count() int { + t.mu.Lock() + defer t.mu.Unlock() + return len(t.entries) +} + +// Clear removes all captured log entries +func (t *testLogger) Clear() { + t.mu.Lock() + defer t.mu.Unlock() + t.entries = make([]LogEntry, 0) +} + +// HasLevel returns true if any log entry with the specified level exists +func (t *testLogger) HasLevel(level slog.Level) bool { + t.mu.Lock() + defer t.mu.Unlock() + for _, entry := range t.entries { + if entry.Level == level { + return true + } + } + return false +} + +// FindByMessage returns all log entries that contain the specified message +func (t *testLogger) FindByMessage(msg string) []LogEntry { + t.mu.Lock() + defer t.mu.Unlock() + result := make([]LogEntry, 0) + for _, entry := range t.entries { + if entry.Message == msg { + result = append(result, entry) + } + } + return result +} diff --git a/logging_test.go b/logging_test.go new file mode 100644 index 0000000..4f5c469 --- /dev/null +++ b/logging_test.go @@ -0,0 +1,209 @@ +package warc + +import ( + "context" + "log/slog" + "testing" +) + +func TestTestLogger_BasicFunctionality(t *testing.T) { + logger := NewTestLogger() + + // Log some messages + logger.Debug("debug message", "key1", "value1") + logger.Info("info message", "key2", "value2") + logger.Warn("warn message", "key3", "value3") + logger.Error("error message", "key4", "value4") + + // Check count + if logger.Count() != 4 { + t.Errorf("Expected 4 log entries, got %d", logger.Count()) + } + + // Test Next() - sequential retrieval + entry := logger.Next() + if entry == nil { + t.Fatal("Expected first entry, got nil") + } + if entry.Level != slog.LevelDebug || entry.Message != "debug message" { + t.Errorf("Expected debug message, got %v: %s", entry.Level, entry.Message) + } + + entry = logger.Next() + if entry == nil { + t.Fatal("Expected second entry, got nil") + } + if entry.Level != slog.LevelInfo || entry.Message != "info message" { + t.Errorf("Expected info message, got %v: %s", entry.Level, entry.Message) + } + + // Count should be 2 now (2 consumed) + if logger.Count() != 2 { + t.Errorf("Expected 2 remaining entries, got %d", logger.Count()) + } +} + +func TestTestLogger_HasLevel(t *testing.T) { + logger := NewTestLogger() + + logger.Debug("test") + logger.Info("test") + + if !logger.HasLevel(slog.LevelDebug) { + t.Error("Expected to find Debug level") + } + + if !logger.HasLevel(slog.LevelInfo) { + t.Error("Expected to find Info level") + } + + if logger.HasLevel(slog.LevelError) { + t.Error("Did not expect to find Error level") + } +} + +func TestTestLogger_FindByMessage(t *testing.T) { + logger := NewTestLogger() + + logger.Info("test message 1") + logger.Info("test message 2") + logger.Info("test message 1") + logger.Debug("other message") + + entries := logger.FindByMessage("test message 1") + if len(entries) != 2 { + t.Errorf("Expected 2 entries with 'test message 1', got %d", len(entries)) + } + + entries = logger.FindByMessage("other message") + if len(entries) != 1 { + t.Errorf("Expected 1 entry with 'other message', got %d", len(entries)) + } + + entries = logger.FindByMessage("nonexistent") + if len(entries) != 0 { + t.Errorf("Expected 0 entries with 'nonexistent', got %d", len(entries)) + } +} + +func TestTestLogger_Clear(t *testing.T) { + logger := NewTestLogger() + + logger.Info("test 1") + logger.Info("test 2") + logger.Info("test 3") + + if logger.Count() != 3 { + t.Errorf("Expected 3 entries before clear, got %d", logger.Count()) + } + + logger.Clear() + + if logger.Count() != 0 { + t.Errorf("Expected 0 entries after clear, got %d", logger.Count()) + } + + // Next should return nil after clear + if entry := logger.Next(); entry != nil { + t.Error("Expected nil after clear, got entry") + } +} + +func TestTestLogger_LogMethod(t *testing.T) { + logger := NewTestLogger() + + ctx := context.Background() + logger.Log(ctx, slog.LevelWarn, "custom level message", "key", "value") + + if logger.Count() != 1 { + t.Errorf("Expected 1 entry, got %d", logger.Count()) + } + + entry := logger.Next() + if entry == nil { + t.Fatal("Expected entry, got nil") + } + if entry.Level != slog.LevelWarn { + t.Errorf("Expected Warn level, got %v", entry.Level) + } + if entry.Message != "custom level message" { + t.Errorf("Expected 'custom level message', got %s", entry.Message) + } +} + +func TestTestLogger_Entries(t *testing.T) { + logger := NewTestLogger() + + logger.Debug("msg1") + logger.Info("msg2") + logger.Error("msg3") + + // Get all entries + entries := logger.Entries() + if len(entries) != 3 { + t.Errorf("Expected 3 entries, got %d", len(entries)) + } + + // Verify it's a copy (modifying shouldn't affect logger) + entries[0].Message = "modified" + allEntries := logger.Entries() + if allEntries[0].Message == "modified" { + t.Error("Entries() should return a copy, but original was modified") + } + + // Count should still be 3 (Entries doesn't consume) + if logger.Count() != 3 { + t.Errorf("Expected count to remain 3, got %d", logger.Count()) + } +} + +func TestTestLogger_ArgsCapture(t *testing.T) { + logger := NewTestLogger() + + logger.Info("test message", "key1", "value1", "key2", 42, "key3", true) + + entry := logger.Next() + if entry == nil { + t.Fatal("Expected entry, got nil") + } + + if len(entry.Args) != 6 { + t.Errorf("Expected 6 args, got %d", len(entry.Args)) + } + + // Verify args are captured correctly + if entry.Args[0] != "key1" || entry.Args[1] != "value1" { + t.Error("Args not captured correctly") + } + if entry.Args[2] != "key2" || entry.Args[3] != 42 { + t.Error("Args not captured correctly") + } + if entry.Args[4] != "key3" || entry.Args[5] != true { + t.Error("Args not captured correctly") + } +} + +func TestTestLogger_ThreadSafety(t *testing.T) { + logger := NewTestLogger() + + // Concurrent writes + done := make(chan bool) + for i := 0; i < 10; i++ { + go func() { + for j := 0; j < 100; j++ { + logger.Info("concurrent message") + } + done <- true + }() + } + + // Wait for all goroutines + for i := 0; i < 10; i++ { + <-done + } + + // Should have 1000 entries + if logger.Count() != 1000 { + t.Errorf("Expected 1000 entries, got %d", logger.Count()) + } +} diff --git a/smoke_test.go b/smoke_test.go index 6411001..bb92844 100644 --- a/smoke_test.go +++ b/smoke_test.go @@ -30,11 +30,11 @@ func TestSmokeWARCFormatRegression(t *testing.T) { // These values were extracted from a known-good WARC file and serve as // a snapshot of correct format behavior. expectedRecords := []struct { - warcType string - contentLength int64 - blockDigest string - payloadDigest string // only for response records - targetURI string // only for response records + warcType string + contentLength int64 + blockDigest string + payloadDigest string // only for response records + targetURI string // only for response records }{ { warcType: "warcinfo", diff --git a/stats.go b/stats.go new file mode 100644 index 0000000..af55ff7 --- /dev/null +++ b/stats.go @@ -0,0 +1,355 @@ +package warc + +import ( + "fmt" + "sort" + "strings" + "sync" + "sync/atomic" +) + +const ( + // totalDataWritten is the name of the metric that tracks the total data written to WARC files. + totalDataWritten string = "total_data_written" + totalDataWrittenHelp string = "Total data written to WARC files in bytes" + + // localDedupedBytesTotal is the name of the metric that tracks the total bytes deduped using local dedupe. + localDedupedBytesTotal string = "local_deduped_bytes_total" + localDedupedBytesTotalHelp string = "Total bytes deduped using local dedupe" + + // localDedupedTotal is the name of the metric that tracks the total records deduped using local dedupe. + localDedupedTotal string = "local_deduped_total" + localDedupedTotalHelp string = "Total records deduped using local dedupe" + + // doppelgangerDedupedBytesTotal is the name of the metric that tracks the total bytes deduped using Doppelganger. + doppelgangerDedupedBytesTotal string = "doppelganger_deduped_bytes_total" + doppelgangerDedupedBytesTotalHelp string = "Total bytes deduped using Doppelganger" + + // doppelgangerDedupedTotal is the name of the metric that tracks the total records deduped using Doppelganger. + doppelgangerDedupedTotal string = "doppelganger_deduped_total" + doppelgangerDedupedTotalHelp string = "Total records deduped using Doppelganger" + + // cdxDedupedBytesTotal is the name of the metric that tracks the total bytes deduped using CDX. + cdxDedupedBytesTotal string = "cdx_deduped_bytes_total" + cdxDedupedBytesTotalHelp string = "Total bytes deduped using CDX" + + // cdxDedupedTotal is the name of the metric that tracks the total records deduped using CDX. + cdxDedupedTotal string = "cdx_deduped_total" + cdxDedupedTotalHelp string = "Total records deduped using CDX" + + // proxyRequestsTotal is the name of the metric that tracks the total number of requests gone through a proxy. + proxyRequestsTotal string = "proxy_requests_total" + proxyRequestsTotalHelp string = "Total number of requests gone through a proxy" + + // proxyErrorsTotal is the name of the metric that tracks the total number of errors occurred with a proxy. + proxyErrorsTotal string = "proxy_errors_total" + proxyErrorsTotalHelp string = "Total number of errors occurred with a proxy" + + // proxyLastUsedNanoseconds is the name of the metric that tracks the last time a proxy was used. + proxyLastUsedNanoseconds string = "proxy_last_used_nanoseconds" + proxyLastUsedNanosecondsHelp string = "Last time a proxy was used in seconds (unix timestamp ns)" +) + +// Labels represents Prometheus-style label values as key-value pairs. +// Labels are used with WithLabels() to specify concrete label values when recording metrics. +// +// Example usage with labels: +// +// registry := newLocalRegistry() +// +// // Register a counter with label names (declares which dimensions the metric tracks) +// httpRequests := registry.RegisterCounter("http_requests_total", "Total HTTP requests", []string{"method", "status"}) +// +// // Use WithLabels to specify label values when recording +// httpRequests.WithLabels(Labels{"method": "GET", "status": "200"}).Inc() +// httpRequests.WithLabels(Labels{"method": "POST", "status": "201"}).Add(5) +// httpRequests.WithLabels(Labels{"method": "GET", "status": "404"}).Inc() +// +// // Each unique combination of label values creates a separate metric series +// // The counter with method=GET,status=200 has value 1 +// // The counter with method=POST,status=201 has value 5 +// // The counter with method=GET,status=404 has value 1 +// +// Example usage without labels: +// +// // For metrics without labels, pass nil or empty slice when registering +// totalCounter := registry.RegisterCounter("requests_total", "All requests", nil) +// // WithLabels(nil) is required but the labels parameter is ignored +// totalCounter.WithLabels(nil).Inc() +// +// Labels are internally sorted by key to ensure consistent metric identification +// regardless of the order in which they are specified. This means: +// +// Labels{"method": "GET", "status": "200"} == Labels{"status": "200", "method": "GET"} +type Labels map[string]string + +// labelsToString converts labels to a consistent string representation for use as map keys. +// Labels are sorted by key to ensure consistent ordering. +func labelsToString(labels Labels) string { + if len(labels) == 0 { + return "" + } + + // Sort keys for consistent ordering + keys := make([]string, 0, len(labels)) + for k := range labels { + keys = append(keys, k) + } + sort.Strings(keys) + + // Build string representation + parts := make([]string, 0, len(labels)) + for _, k := range keys { + parts = append(parts, fmt.Sprintf("%s=%q", k, labels[k])) + } + return strings.Join(parts, ",") +} + +// makeMetricKey creates a unique key for a metric with labels. +func makeMetricKey(name string, labels Labels) string { + if len(labels) == 0 { + return name + } + return name + "{" + labelsToString(labels) + "}" +} + +// RegistryOpts is an interface that provides a WithLabels method to get a metric with specific labels. +type RegistryOpts[T any] interface { + // WithLabels returns a Counter for the given label values. + // If the counter has no labels, labels parameter is ignored. + WithLabels(labels Labels) T +} + +// Counter represents a monotonically increasing metric. +type Counter interface { + RegistryOpts[Counter] + // Inc increments the counter by 1. + Inc() + // Add adds the given value to the counter. + Add(value int64) + // Get returns the current value of the counter. + // This is used to support unit-testing of the metrics. + Get() int64 +} + +// Gauge represents a metric that can go up or down. +type Gauge interface { + RegistryOpts[Gauge] + // Set sets the gauge to the given value. + Set(value int64) + // Inc increments the gauge by 1. + Inc() + // Dec decrements the gauge by 1. + Dec() + // Add adds the given value to the gauge. + Add(value int64) + // Sub subtracts the given value from the gauge. + Sub(value int64) + // Get returns the current value of the gauge. + // This is used to support unit-testing of the metrics. + Get() int64 +} + +// Histogram represents a metric for observing distributions of values. +type Histogram interface { + RegistryOpts[Histogram] + // Observe adds a single observation to the histogram. + Observe(value int64) +} + +// StatsRegistry provides a registry for external libraries to register and update metrics. +// The StatsRegistry implementation is expected to be thread-safe so that gowarc can safely register and update metrics from multiple goroutines. +type StatsRegistry interface { + // RegisterCounter registers a new counter metric with optional label names. + // labelNames specifies which labels this counter will use (e.g., []string{"method", "status"}). + // Returns a Counter that can be used with WithLabels() to specify label values. + // If labelNames is nil or empty, returns a Counter that ignores labels. + RegisterCounter(name, help string, labelNames []string) Counter + + // RegisterGauge registers a new gauge metric with optional label names. + // labelNames specifies which labels this gauge will use (e.g., []string{"location", "type"}). + // Returns a Gauge that can be used with WithLabels() to specify label values. + // If labelNames is nil or empty, returns a Gauge that ignores labels. + RegisterGauge(name, help string, labelNames []string) Gauge + + // RegisterHistogram registers a new histogram metric with the given buckets and optional label names. + // If buckets is nil, uses Prometheus default buckets. + // labelNames specifies which labels this histogram will use (e.g., []string{"endpoint", "method"}). + // Returns a Histogram that can be used with WithLabels() to specify label values. + // If labelNames is nil or empty, returns a Histogram that ignores labels. + RegisterHistogram(name, help string, buckets []int64, labelNames []string) Histogram +} + +// Nil-safe implementations for when no StatsRegistry is provided. +type localCounter struct { + v atomic.Int64 +} + +func (n *localCounter) WithLabels(_ Labels) Counter { return n } +func (n *localCounter) Inc() { n.v.Add(1) } +func (n *localCounter) Add(value int64) { n.v.Add(value) } +func (n *localCounter) Get() int64 { return n.v.Load() } + +type localGauge struct { + v atomic.Int64 +} + +func (n *localGauge) WithLabels(_ Labels) Gauge { return n } +func (n *localGauge) Set(value int64) { n.v.Store(value) } +func (n *localGauge) Inc() { n.v.Add(1) } +func (n *localGauge) Dec() { n.v.Add(-1) } +func (n *localGauge) Add(value int64) { n.v.Add(value) } +func (n *localGauge) Sub(value int64) { n.v.Add(-value) } +func (n *localGauge) Get() int64 { return n.v.Load() } + +type localHistogram struct{} + +func (n *localHistogram) WithLabels(_ Labels) Histogram { return n } +func (n *localHistogram) Observe(_ int64) {} + +// labeledCounter implements Counter with label support +type labeledCounter struct { + name string + registry *localRegistry +} + +func (l *labeledCounter) WithLabels(labels Labels) Counter { + return l.registry.getOrCreateCounter(l.name, labels) +} + +func (l *labeledCounter) Inc() { + l.registry.getOrCreateCounter(l.name, nil).Inc() +} + +func (l *labeledCounter) Add(value int64) { + l.registry.getOrCreateCounter(l.name, nil).Add(value) +} + +func (l *labeledCounter) Get() int64 { + return l.registry.getOrCreateCounter(l.name, nil).Get() +} + +// labeledGauge implements Gauge with label support +type labeledGauge struct { + name string + registry *localRegistry +} + +func (l *labeledGauge) WithLabels(labels Labels) Gauge { + return l.registry.getOrCreateGauge(l.name, labels) +} + +func (l *labeledGauge) Set(value int64) { + l.registry.getOrCreateGauge(l.name, nil).Set(value) +} + +func (l *labeledGauge) Inc() { + l.registry.getOrCreateGauge(l.name, nil).Inc() +} + +func (l *labeledGauge) Dec() { + l.registry.getOrCreateGauge(l.name, nil).Dec() +} + +func (l *labeledGauge) Add(value int64) { + l.registry.getOrCreateGauge(l.name, nil).Add(value) +} + +func (l *labeledGauge) Sub(value int64) { + l.registry.getOrCreateGauge(l.name, nil).Sub(value) +} + +func (l *labeledGauge) Get() int64 { + return l.registry.getOrCreateGauge(l.name, nil).Get() +} + +// labeledHistogram implements Histogram with label support +type labeledHistogram struct { + name string + registry *localRegistry +} + +func (l *labeledHistogram) WithLabels(labels Labels) Histogram { + return l.registry.getOrCreateHistogram(l.name, labels) +} + +func (l *labeledHistogram) Observe(value int64) { + l.registry.getOrCreateHistogram(l.name, nil).Observe(value) +} + +type localRegistry struct { + sync.Mutex + gauges map[string]*localGauge + counters map[string]*localCounter + histograms map[string]*localHistogram +} + +func newLocalRegistry() *localRegistry { + return &localRegistry{} +} + +func (n *localRegistry) getOrCreateCounter(name string, labels Labels) Counter { + n.Lock() + defer n.Unlock() + if n.counters == nil { + n.counters = make(map[string]*localCounter) + } + key := makeMetricKey(name, labels) + if c, ok := n.counters[key]; ok { + return c + } + c := &localCounter{} + n.counters[key] = c + return c +} + +func (n *localRegistry) getOrCreateGauge(name string, labels Labels) Gauge { + n.Lock() + defer n.Unlock() + if n.gauges == nil { + n.gauges = make(map[string]*localGauge) + } + key := makeMetricKey(name, labels) + if g, ok := n.gauges[key]; ok { + return g + } + g := &localGauge{} + n.gauges[key] = g + return g +} + +func (n *localRegistry) getOrCreateHistogram(name string, labels Labels) Histogram { + n.Lock() + defer n.Unlock() + if n.histograms == nil { + n.histograms = make(map[string]*localHistogram) + } + key := makeMetricKey(name, labels) + if h, ok := n.histograms[key]; ok { + return h + } + h := &localHistogram{} + n.histograms[key] = h + return h +} + +func (n *localRegistry) RegisterCounter(name, _ string, _ []string) Counter { + return &labeledCounter{ + name: name, + registry: n, + } +} + +func (n *localRegistry) RegisterGauge(name, _ string, _ []string) Gauge { + return &labeledGauge{ + name: name, + registry: n, + } +} + +func (n *localRegistry) RegisterHistogram(name, _ string, _ []int64, _ []string) Histogram { + return &labeledHistogram{ + name: name, + registry: n, + } +} diff --git a/stats_test.go b/stats_test.go new file mode 100644 index 0000000..bb9f04d --- /dev/null +++ b/stats_test.go @@ -0,0 +1,607 @@ +package warc + +import ( + "sync" + "testing" +) + +// TestLocalCounter tests the localCounter implementation +func TestLocalCounter(t *testing.T) { + c := &localCounter{} + + // Test initial value + if c.Get() != 0 { + t.Errorf("Expected initial value 0, got %d", c.Get()) + } + + // Test Inc + c.Inc() + if c.Get() != 1 { + t.Errorf("Expected value 1 after Inc, got %d", c.Get()) + } + + // Test Add + c.Add(5) + if c.Get() != 6 { + t.Errorf("Expected value 6 after Add(5), got %d", c.Get()) + } + + // Test Add with negative value (counters should still accept it) + c.Add(-2) + if c.Get() != 4 { + t.Errorf("Expected value 4 after Add(-2), got %d", c.Get()) + } +} + +// TestLocalGauge tests the localGauge implementation +func TestLocalGauge(t *testing.T) { + g := &localGauge{} + + // Test initial value + if g.Get() != 0 { + t.Errorf("Expected initial value 0, got %d", g.Get()) + } + + // Test Set + g.Set(10) + if g.Get() != 10 { + t.Errorf("Expected value 10 after Set(10), got %d", g.Get()) + } + + // Test Inc + g.Inc() + if g.Get() != 11 { + t.Errorf("Expected value 11 after Inc, got %d", g.Get()) + } + + // Test Dec + g.Dec() + if g.Get() != 10 { + t.Errorf("Expected value 10 after Dec, got %d", g.Get()) + } + + // Test Add + g.Add(5) + if g.Get() != 15 { + t.Errorf("Expected value 15 after Add(5), got %d", g.Get()) + } + + // Test Sub + g.Sub(3) + if g.Get() != 12 { + t.Errorf("Expected value 12 after Sub(3), got %d", g.Get()) + } + + // Test Set to negative value + g.Set(-5) + if g.Get() != -5 { + t.Errorf("Expected value -5 after Set(-5), got %d", g.Get()) + } +} + +// TestLocalHistogram tests the localHistogram implementation +func TestLocalHistogram(t *testing.T) { + h := &localHistogram{} + + // Histogram's Observe method is a no-op, just ensure it doesn't panic + h.Observe(100) + h.Observe(0) + h.Observe(-50) +} + +// TestLocalRegistryRegisterCounter tests the RegisterCounter method +func TestLocalRegistryRegisterCounter(t *testing.T) { + registry := newLocalRegistry() + + // Register a new counter + counter1 := registry.RegisterCounter("test_counter", "Test counter help", nil).WithLabels(nil) + if counter1 == nil { + t.Fatal("Expected counter to be created, got nil") + } + + // Verify it's in the registry + if len(registry.counters) != 1 { + t.Errorf("Expected 1 counter in registry, got %d", len(registry.counters)) + } + + // Register the same counter again - should return existing one + counter2 := registry.RegisterCounter("test_counter", "Different help text", nil).WithLabels(nil) + if counter1 != counter2 { + t.Error("Expected RegisterCounter to return existing counter for same name") + } + + // Verify still only one counter + if len(registry.counters) != 1 { + t.Errorf("Expected 1 counter in registry after re-registration, got %d", len(registry.counters)) + } + + // Register a different counter + counter3 := registry.RegisterCounter("another_counter", "Another counter", nil).WithLabels(nil) + if counter1 == counter3 { + t.Error("Expected different counter instances for different names") + } + + if len(registry.counters) != 2 { + t.Errorf("Expected 2 counters in registry, got %d", len(registry.counters)) + } +} + +// TestLocalRegistryRegisterGauge tests the RegisterGauge method +func TestLocalRegistryRegisterGauge(t *testing.T) { + registry := newLocalRegistry() + + // Register a new gauge + gauge1 := registry.RegisterGauge("test_gauge", "Test gauge help", nil).WithLabels(nil) + if gauge1 == nil { + t.Fatal("Expected gauge to be created, got nil") + } + + // Verify it's in the registry + if len(registry.gauges) != 1 { + t.Errorf("Expected 1 gauge in registry, got %d", len(registry.gauges)) + } + + // Register the same gauge again - should return existing one + gauge2 := registry.RegisterGauge("test_gauge", "Different help text", nil).WithLabels(nil) + if gauge1 != gauge2 { + t.Error("Expected RegisterGauge to return existing gauge for same name") + } + + // Verify still only one gauge + if len(registry.gauges) != 1 { + t.Errorf("Expected 1 gauge in registry after re-registration, got %d", len(registry.gauges)) + } + + // Register a different gauge + gauge3 := registry.RegisterGauge("another_gauge", "Another gauge", nil).WithLabels(nil) + if gauge1 == gauge3 { + t.Error("Expected different gauge instances for different names") + } + + if len(registry.gauges) != 2 { + t.Errorf("Expected 2 gauges in registry, got %d", len(registry.gauges)) + } +} + +// TestLocalRegistryRegisterHistogram tests the RegisterHistogram method +func TestLocalRegistryRegisterHistogram(t *testing.T) { + registry := newLocalRegistry() + + // Register a new histogram + labeledHist1 := registry.RegisterHistogram("test_histogram", "Test histogram help", []int64{1, 2, 3}, nil) + if labeledHist1 == nil { + t.Fatal("Expected labeled histogram to be created, got nil") + } + + // Get the actual histogram instance with WithLabels + histogram1 := labeledHist1.WithLabels(nil) + if histogram1 == nil { + t.Fatal("Expected histogram to be created, got nil") + } + + // Verify it's in the registry + if len(registry.histograms) != 1 { + t.Errorf("Expected 1 histogram in registry, got %d", len(registry.histograms)) + } + + // Register the same histogram again - calling WithLabels should return the same instance + histogram2 := registry.RegisterHistogram("test_histogram", "Different help text", []int64{5, 10, 15}, nil).WithLabels(nil) + if histogram1 != histogram2 { + t.Error("Expected same histogram instance for same name and labels") + } + + // Verify still only one histogram + if len(registry.histograms) != 1 { + t.Errorf("Expected 1 histogram in registry after re-registration, got %d", len(registry.histograms)) + } + + // Register and use a different histogram + registry.RegisterHistogram("another_histogram", "Another histogram", nil, nil).WithLabels(nil) + + if len(registry.histograms) != 2 { + t.Errorf("Expected 2 histograms in registry, got %d", len(registry.histograms)) + } + + // Verify both histograms are in the registry by key + // Note: keys are now the full metric keys including label info + if len(registry.histograms) < 2 { + t.Error("Expected both histograms to be in registry") + } +} + +// TestLocalRegistryCounterFunctionality tests that registered counters work correctly +func TestLocalRegistryCounterFunctionality(t *testing.T) { + registry := newLocalRegistry() + counter := registry.RegisterCounter("functional_counter", "Test", nil).WithLabels(nil) + + counter.Inc() + if counter.Get() != 1 { + t.Errorf("Expected counter value 1, got %d", counter.Get()) + } + + counter.Add(10) + if counter.Get() != 11 { + t.Errorf("Expected counter value 11, got %d", counter.Get()) + } +} + +// TestLocalRegistryGaugeFunctionality tests that registered gauges work correctly +func TestLocalRegistryGaugeFunctionality(t *testing.T) { + registry := newLocalRegistry() + gauge := registry.RegisterGauge("functional_gauge", "Test", nil).WithLabels(nil) + + gauge.Set(100) + if gauge.Get() != 100 { + t.Errorf("Expected gauge value 100, got %d", gauge.Get()) + } + + gauge.Inc() + if gauge.Get() != 101 { + t.Errorf("Expected gauge value 101, got %d", gauge.Get()) + } + + gauge.Dec() + if gauge.Get() != 100 { + t.Errorf("Expected gauge value 100, got %d", gauge.Get()) + } +} + +// TestLocalRegistryConcurrentAccess tests thread-safety of the localRegistry +func TestLocalRegistryConcurrentAccess(t *testing.T) { + registry := newLocalRegistry() + var wg sync.WaitGroup + + // Number of concurrent goroutines + numGoroutines := 100 + + // Test concurrent counter registration + wg.Add(numGoroutines) + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + // All goroutines try to register the same counter + counter := registry.RegisterCounter("shared_counter", "Test", nil).WithLabels(nil) + counter.Inc() + }(i) + } + wg.Wait() + + // Should have exactly one counter + if len(registry.counters) != 1 { + t.Errorf("Expected 1 counter after concurrent registration, got %d", len(registry.counters)) + } + + // Counter should have been incremented by all goroutines + counter := registry.RegisterCounter("shared_counter", "Test", nil).WithLabels(nil) + if counter.Get() != int64(numGoroutines) { + t.Errorf("Expected counter value %d, got %d", numGoroutines, counter.Get()) + } + + // Test concurrent gauge registration + wg.Add(numGoroutines) + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + gauge := registry.RegisterGauge("shared_gauge", "Test", nil).WithLabels(nil) + gauge.Inc() + }(i) + } + wg.Wait() + + // Should have exactly one gauge + if len(registry.gauges) != 1 { + t.Errorf("Expected 1 gauge after concurrent registration, got %d", len(registry.gauges)) + } + + // Test concurrent histogram registration + wg.Add(numGoroutines) + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + histogram := registry.RegisterHistogram("shared_histogram", "Test", nil, nil).WithLabels(nil) + histogram.Observe(int64(id)) + }(i) + } + wg.Wait() + + // Should have exactly one histogram + if len(registry.histograms) != 1 { + t.Errorf("Expected 1 histogram after concurrent registration, got %d", len(registry.histograms)) + } +} + +// TestLocalRegistryMultipleMetrics tests registering multiple different metrics +func TestLocalRegistryMultipleMetrics(t *testing.T) { + registry := newLocalRegistry() + + // Register multiple counters + for i := 0; i < 5; i++ { + name := "counter_" + string(rune('a'+i)) + registry.RegisterCounter(name, "Test counter", nil).WithLabels(nil) + } + + // Register multiple gauges + for i := 0; i < 5; i++ { + name := "gauge_" + string(rune('a'+i)) + registry.RegisterGauge(name, "Test gauge", nil).WithLabels(nil) + } + + // Register multiple histograms + for i := 0; i < 5; i++ { + name := "histogram_" + string(rune('a'+i)) + registry.RegisterHistogram(name, "Test histogram", nil, nil).WithLabels(nil) + } + + if len(registry.counters) != 5 { + t.Errorf("Expected 5 counters, got %d", len(registry.counters)) + } + + if len(registry.gauges) != 5 { + t.Errorf("Expected 5 gauges, got %d", len(registry.gauges)) + } + + if len(registry.histograms) != 5 { + t.Errorf("Expected 5 histograms, got %d", len(registry.histograms)) + } +} + +// TestLocalCounterConcurrentIncrement tests concurrent increments on a counter +func TestLocalCounterConcurrentIncrement(t *testing.T) { + counter := &localCounter{} + var wg sync.WaitGroup + numGoroutines := 1000 + incrementsPerGoroutine := 100 + + wg.Add(numGoroutines) + for i := 0; i < numGoroutines; i++ { + go func() { + defer wg.Done() + for j := 0; j < incrementsPerGoroutine; j++ { + counter.Inc() + } + }() + } + wg.Wait() + + expected := int64(numGoroutines * incrementsPerGoroutine) + if counter.Get() != expected { + t.Errorf("Expected counter value %d, got %d", expected, counter.Get()) + } +} + +// TestLocalGaugeConcurrentOperations tests concurrent operations on a gauge +func TestLocalGaugeConcurrentOperations(t *testing.T) { + gauge := &localGauge{} + var wg sync.WaitGroup + numGoroutines := 100 + + // Set initial value + gauge.Set(0) + + // Half goroutines increment, half decrement + wg.Add(numGoroutines) + for i := 0; i < numGoroutines; i++ { + if i%2 == 0 { + go func() { + defer wg.Done() + gauge.Inc() + }() + } else { + go func() { + defer wg.Done() + gauge.Dec() + }() + } + } + wg.Wait() + + // With equal increments and decrements, value should be 0 + if gauge.Get() != 0 { + t.Errorf("Expected gauge value 0 after equal Inc/Dec operations, got %d", gauge.Get()) + } +} + +// TestLabelsToString tests the labelsToString function +func TestLabelsToString(t *testing.T) { + tests := []struct { + name string + labels Labels + expected string + }{ + { + name: "empty labels", + labels: nil, + expected: "", + }, + { + name: "single label", + labels: Labels{"method": "GET"}, + expected: `method="GET"`, + }, + { + name: "multiple labels sorted", + labels: Labels{"method": "GET", "status": "200"}, + expected: `method="GET",status="200"`, + }, + { + name: "multiple labels reverse order", + labels: Labels{"status": "200", "method": "GET"}, + expected: `method="GET",status="200"`, // Should be sorted + }, + { + name: "labels with special characters", + labels: Labels{"path": "/api/v1/users", "method": "POST"}, + expected: `method="POST",path="/api/v1/users"`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := labelsToString(tt.labels) + if result != tt.expected { + t.Errorf("labelsToString() = %q, expected %q", result, tt.expected) + } + }) + } +} + +// TestMakeMetricKey tests the makeMetricKey function +func TestMakeMetricKey(t *testing.T) { + tests := []struct { + name string + metricName string + labels Labels + expected string + }{ + { + name: "no labels", + metricName: "http_requests_total", + labels: nil, + expected: "http_requests_total", + }, + { + name: "with labels", + metricName: "http_requests_total", + labels: Labels{"method": "GET", "status": "200"}, + expected: `http_requests_total{method="GET",status="200"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := makeMetricKey(tt.metricName, tt.labels) + if result != tt.expected { + t.Errorf("makeMetricKey() = %q, expected %q", result, tt.expected) + } + }) + } +} + +// TestLocalRegistryWithLabels tests that metrics with different labels are separate +func TestLocalRegistryWithLabels(t *testing.T) { + registry := newLocalRegistry() + + // Register a counter with label names + httpRequests := registry.RegisterCounter("http_requests", "HTTP requests", []string{"method"}) + + // Use WithLabels to get counters for different label values + counter1 := httpRequests.WithLabels(Labels{"method": "GET"}) + counter2 := httpRequests.WithLabels(Labels{"method": "POST"}) + + // Register a counter without labels + totalRequests := registry.RegisterCounter("total_requests", "Total requests", nil) + counter3 := totalRequests.WithLabels(nil) + + // Increment each counter by different amounts + counter1.Add(10) + counter2.Add(20) + counter3.Add(30) + + // Verify values are independent + if counter1.Get() != 10 { + t.Errorf("Expected counter1 value 10, got %d", counter1.Get()) + } + if counter2.Get() != 20 { + t.Errorf("Expected counter2 value 20, got %d", counter2.Get()) + } + if counter3.Get() != 30 { + t.Errorf("Expected counter3 value 30, got %d", counter3.Get()) + } + + // Verify all three are in the registry + if len(registry.counters) != 3 { + t.Errorf("Expected 3 counters in registry, got %d", len(registry.counters)) + } +} + +// TestLocalRegistryGaugeWithLabels tests gauges with labels +func TestLocalRegistryGaugeWithLabels(t *testing.T) { + registry := newLocalRegistry() + + // Register a gauge with label names + temperature := registry.RegisterGauge("temperature", "Temperature", []string{"location"}) + + // Use WithLabels to get gauges for different label values + gauge1 := temperature.WithLabels(Labels{"location": "indoor"}) + gauge2 := temperature.WithLabels(Labels{"location": "outdoor"}) + + // Set different values + gauge1.Set(20) + gauge2.Set(15) + + // Verify values are independent + if gauge1.Get() != 20 { + t.Errorf("Expected gauge1 value 20, got %d", gauge1.Get()) + } + if gauge2.Get() != 15 { + t.Errorf("Expected gauge2 value 15, got %d", gauge2.Get()) + } + + // Verify both are in the registry + if len(registry.gauges) != 2 { + t.Errorf("Expected 2 gauges in registry, got %d", len(registry.gauges)) + } +} + +// TestLocalRegistryHistogramWithLabels tests histograms with labels +func TestLocalRegistryHistogramWithLabels(t *testing.T) { + registry := newLocalRegistry() + + // Register a histogram with label names + responseTime := registry.RegisterHistogram("response_time", "Response time", []int64{100, 200, 500}, []string{"endpoint"}) + + // Use WithLabels to get histograms for different label values + hist1 := responseTime.WithLabels(Labels{"endpoint": "/api"}) + hist2 := responseTime.WithLabels(Labels{"endpoint": "/health"}) + + // Observe values (no-ops for local implementation, but verifies no panic) + hist1.Observe(150) + hist2.Observe(50) + + // Verify both are in the registry as separate entries + if len(registry.histograms) != 2 { + t.Errorf("Expected 2 histograms in registry, got %d", len(registry.histograms)) + } + + // Verify we can retrieve them again with the same labels + hist1Again := responseTime.WithLabels(Labels{"endpoint": "/api"}) + hist2Again := responseTime.WithLabels(Labels{"endpoint": "/health"}) + + // Should get the same instances back + if hist1 != hist1Again { + t.Error("Expected to get same histogram instance for same labels") + } + if hist2 != hist2Again { + t.Error("Expected to get same histogram instance for same labels") + } +} + +// TestLocalRegistrySameLabelsDifferentOrder tests that label order doesn't matter +func TestLocalRegistrySameLabelsDifferentOrder(t *testing.T) { + registry := newLocalRegistry() + + // Register counter with label names + requests := registry.RegisterCounter("requests", "Requests", []string{"method", "status"}) + + // Get counter with labels in one order + counter1 := requests.WithLabels(Labels{"method": "GET", "status": "200"}) + counter1.Add(5) + + // Get counter with same labels in different order + counter2 := requests.WithLabels(Labels{"status": "200", "method": "GET"}) + + // Should return the same counter + if counter1 != counter2 { + t.Error("Expected same counter for same labels in different order") + } + + // Value should be preserved + if counter2.Get() != 5 { + t.Errorf("Expected counter value 5, got %d", counter2.Get()) + } + + // Should only have one counter in registry + if len(registry.counters) != 1 { + t.Errorf("Expected 1 counter in registry, got %d", len(registry.counters)) + } +} diff --git a/utils.go b/utils.go index e053f8f..df6386c 100644 --- a/utils.go +++ b/utils.go @@ -6,6 +6,7 @@ import ( "encoding/binary" "errors" "io" + "net/url" "os" "strings" "time" @@ -40,7 +41,7 @@ func isHTTPRequest(line string) bool { } // NewWriter creates a new WARC writer. -func NewWriter(writer io.Writer, fileName string, digestAlgorithm DigestAlgorithm, compression string, contentLengthHeader string, newFileCreation bool, dictionary []byte) (*Writer, error) { +func NewWriter(writer io.Writer, fileName string, digestAlgorithm DigestAlgorithm, compression string, contentLengthHeader string, newFileCreation bool, dictionary []byte, stats StatsRegistry, logBackend LogBackend) (*Writer, error) { if compression != "" { switch strings.ToLower(compression) { case "gzip": @@ -52,6 +53,8 @@ func NewWriter(writer io.Writer, fileName string, digestAlgorithm DigestAlgorith DigestAlgorithm: digestAlgorithm, GZIPWriter: gzipWriter, FileWriter: bufio.NewWriter(gzipWriter), + stats: stats, + logBackend: logBackend, }, nil case "zstd": if newFileCreation && len(dictionary) > 0 { @@ -94,6 +97,8 @@ func NewWriter(writer io.Writer, fileName string, digestAlgorithm DigestAlgorith DigestAlgorithm: digestAlgorithm, ZSTDWriter: zstdWriter, FileWriter: bufio.NewWriter(zstdWriter), + stats: stats, + logBackend: logBackend, }, nil } else { zstdWriter, err := zstd.NewWriter(writer, zstd.WithEncoderLevel(zstd.SpeedBetterCompression)) @@ -106,6 +111,8 @@ func NewWriter(writer io.Writer, fileName string, digestAlgorithm DigestAlgorith DigestAlgorithm: digestAlgorithm, ZSTDWriter: zstdWriter, FileWriter: bufio.NewWriter(zstdWriter), + stats: stats, + logBackend: logBackend, }, nil } default: @@ -118,6 +125,8 @@ func NewWriter(writer io.Writer, fileName string, digestAlgorithm DigestAlgorith Compression: "", DigestAlgorithm: digestAlgorithm, FileWriter: bufio.NewWriter(writer), + stats: stats, + logBackend: logBackend, }, nil } @@ -227,3 +236,17 @@ func getContentLength(rwsc spooledtempfile.ReadWriteSeekCloser) int { return int(fileInfo.Size()) } } + +func proxyName(u *url.URL) string { + // get domain and replace dots and colons with underscores + domain := strings.ReplaceAll(u.Hostname(), ".", "_") + domain = strings.ReplaceAll(domain, ":", "_") + // get port and replace colons with underscores + port := strings.ReplaceAll(u.Port(), ":", "_") + // if port is empty, set it to 80 + if port == "" { + port = "80" + } + // return domain and port + return domain + "_" + port +} diff --git a/utils_test.go b/utils_test.go index 8903d6f..5b09ffb 100644 --- a/utils_test.go +++ b/utils_test.go @@ -1,6 +1,7 @@ package warc import ( + "net/url" "testing" ) @@ -47,3 +48,92 @@ func TestIsHTTPRequest(t *testing.T) { } } } + +// Tests for the proxyName function +func TestProxyName(t *testing.T) { + tests := []struct { + name string + urlStr string + expected string + }{ + { + name: "domain with explicit port", + urlStr: "http://example.com:8080", + expected: "example_com_8080", + }, + { + name: "domain without port (http defaults to empty, should use 80)", + urlStr: "http://example.com", + expected: "example_com_80", + }, + { + name: "domain without port (https defaults to empty, should use 80)", + urlStr: "https://example.com", + expected: "example_com_80", + }, + { + name: "domain with subdomain and port", + urlStr: "http://api.example.com:3000", + expected: "api_example_com_3000", + }, + { + name: "domain with subdomain without port", + urlStr: "http://api.example.com", + expected: "api_example_com_80", + }, + { + name: "localhost with port", + urlStr: "http://localhost:8080", + expected: "localhost_8080", + }, + { + name: "localhost without port", + urlStr: "http://localhost", + expected: "localhost_80", + }, + { + name: "IPv4 address with port", + urlStr: "http://192.168.1.1:8080", + expected: "192_168_1_1_8080", + }, + { + name: "IPv4 address without port", + urlStr: "http://192.168.1.1", + expected: "192_168_1_1_80", + }, + { + name: "IPv6 address with port", + urlStr: "http://[2001:db8::1]:8080", + expected: "2001_db8__1_8080", + }, + { + name: "IPv6 address without port", + urlStr: "http://[2001:db8::1]", + expected: "2001_db8__1_80", + }, + { + name: "domain with port 443", + urlStr: "https://example.com:443", + expected: "example_com_443", + }, + { + name: "domain with multiple subdomains", + urlStr: "http://api.v2.example.com:9000", + expected: "api_v2_example_com_9000", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + u, err := url.Parse(tt.urlStr) + if err != nil { + t.Fatalf("Failed to parse URL %s: %v", tt.urlStr, err) + } + + result := proxyName(u) + if result != tt.expected { + t.Errorf("proxyName(%s) = %s, expected %s", tt.urlStr, result, tt.expected) + } + }) + } +} diff --git a/warc.go b/warc.go index d8e111a..8987aa5 100644 --- a/warc.go +++ b/warc.go @@ -32,21 +32,12 @@ type RotatorSettings struct { WARCSize float64 // WARCWriterPoolSize defines the number of parallel WARC writers WARCWriterPoolSize int + // StatsRegistry is used to store stats about gowarc + StatsRegistry StatsRegistry + // LogBackend is used to log events from gowarc + LogBackend LogBackend } -var ( - // Create a couple of counters for tracking various stats - DataTotal atomic.Int64 - - CDXDedupeTotalBytes atomic.Int64 - DoppelgangerDedupeTotalBytes atomic.Int64 - LocalDedupeTotalBytes atomic.Int64 - - CDXDedupeTotal atomic.Int64 - DoppelgangerDedupeTotal atomic.Int64 - LocalDedupeTotal atomic.Int64 -) - // NewWARCRotator creates and return a channel that can be used // to communicate records to be written to WARC files to the // recordWriter function running in a goroutine @@ -113,13 +104,14 @@ func recordWriter(settings *RotatorSettings, records chan *RecordBatch, done cha ) // Create and open the initial file + settings.LogBackend.Info("creating initial WARC file", "file", currentFileName) warcFile, err := os.Create(settings.OutputDirectory + currentFileName) if err != nil { panic(err) } // Initialize WARC writer - warcWriter, err := NewWriter(warcFile, currentFileName, settings.digestAlgorithm, settings.Compression, "", true, dictionary) + warcWriter, err := NewWriter(warcFile, currentFileName, settings.digestAlgorithm, settings.Compression, "", true, dictionary, settings.StatsRegistry, settings.LogBackend) if err != nil { panic(err) } @@ -137,7 +129,7 @@ func recordWriter(settings *RotatorSettings, records chan *RecordBatch, done cha panic(err) } - warcWriter, err = NewWriter(warcFile, currentFileName, settings.digestAlgorithm, settings.Compression, "", false, dictionary) + warcWriter, err = NewWriter(warcFile, currentFileName, settings.digestAlgorithm, settings.Compression, "", false, dictionary, settings.StatsRegistry, settings.LogBackend) if err != nil { panic(err) } @@ -148,6 +140,7 @@ func recordWriter(settings *RotatorSettings, records chan *RecordBatch, done cha if more { if isFileSizeExceeded(warcFile, settings.WARCSize) { // WARC file size exceeded settings.WarcSize + settings.LogBackend.Info("WARC file size limit exceeded, rotating to new file", "currentFile", currentFileName, "sizeLimit", settings.WARCSize) // The WARC file is renamed to remove the .open suffix err := os.Rename(path.Join(settings.OutputDirectory, currentFileName), strings.TrimSuffix(path.Join(settings.OutputDirectory, currentFileName), ".open")) if err != nil { @@ -170,13 +163,14 @@ func recordWriter(settings *RotatorSettings, records chan *RecordBatch, done cha // Create the new file and automatically increment the serial inside of GenerateWarcFileName currentFileName = getNextWARCFilename(settings.OutputDirectory, settings.Prefix, settings.Compression, serial) + settings.LogBackend.Info("creating new WARC file after rotation", "file", currentFileName) warcFile, err = os.Create(settings.OutputDirectory + currentFileName) if err != nil { panic(err) } // Initialize new WARC writer - warcWriter, err = NewWriter(warcFile, currentFileName, settings.digestAlgorithm, settings.Compression, "", true, dictionary) + warcWriter, err = NewWriter(warcFile, currentFileName, settings.digestAlgorithm, settings.Compression, "", true, dictionary, settings.StatsRegistry, settings.LogBackend) if err != nil { panic(err) } @@ -198,7 +192,7 @@ func recordWriter(settings *RotatorSettings, records chan *RecordBatch, done cha // Write all the records of the record batch for _, record := range recordBatch.Records { - warcWriter, err = NewWriter(warcFile, currentFileName, settings.digestAlgorithm, settings.Compression, record.Header.Get("Content-Length"), false, dictionary) + warcWriter, err = NewWriter(warcFile, currentFileName, settings.digestAlgorithm, settings.Compression, record.Header.Get("Content-Length"), false, dictionary, settings.StatsRegistry, settings.LogBackend) if err != nil { panic(err) } @@ -251,6 +245,7 @@ func recordWriter(settings *RotatorSettings, records chan *RecordBatch, done cha panic(err) } + settings.LogBackend.Info("WARC writer shutting down cleanly", "finalFile", currentFileName) done <- true return diff --git a/write.go b/write.go index a173856..6088f33 100644 --- a/write.go +++ b/write.go @@ -8,8 +8,8 @@ import ( "strings" "time" - "github.com/internetarchive/gowarc/pkg/spooledtempfile" "github.com/google/uuid" + "github.com/internetarchive/gowarc/pkg/spooledtempfile" "github.com/klauspost/compress/zstd" ) @@ -22,6 +22,9 @@ type Writer struct { Compression string DigestAlgorithm DigestAlgorithm ParallelGZIP bool + + stats StatsRegistry + logBackend LogBackend } // RecordBatch is a structure that contains a bunch of @@ -104,11 +107,13 @@ func (w *Writer) WriteRecord(r *Record) (recordID string, err error) { r.Content.Seek(0, 0) if written, err = io.Copy(w.FileWriter, r.Content); err != nil { + w.logBackend.Error("failed to write WARC record content", "file", w.FileName, "error", err) return recordID, err } if written > 0 { - DataTotal.Add(written) + w.stats.RegisterCounter(totalDataWritten, totalDataWrittenHelp, nil).WithLabels(nil).Add(written) + w.logBackend.Debug("WARC record written", "file", w.FileName, "bytes", written, "recordID", recordID) } if _, err := io.WriteString(w.FileWriter, "\r\n\r\n"); err != nil {