feat: add explicit exchange commit and discard

Author: yzqzss

README.md

@@ -72,11 +72,14 @@ func main() {
 	if err != nil {
 		panic(err)
 	}
+	// Releases the capture on every early-return path. Once Commit succeeds,
+	// this deferred Discard cannot change the decision.
+	defer exchange.Discard(context.Background())
 	// Process response
 	_, _ = io.Copy(io.Discard, exchange.Response.Body)
 	_ = exchange.Response.Body.Close()
-	// Will block until records are actually written to the WARC file
-	if _, err := exchange.Wait(context.Background()); err != nil {
+	// Keep the exchange and wait until its records reach the WARC writer.
+	if _, err := exchange.Commit(context.Background()); err != nil {
 		panic(err)
 	}
 	finalized, err := client.Shutdown(context.Background())
@@ -91,7 +94,9 @@ HTTP/1 captures contain the plaintext HTTP/1 wire bytes seen by the transport.
 HTTP/2 and HTTP/3 captures are deterministic `application/http`
 serializations of the actual stream headers, body data, and trailers.
 Closing a response body early performs a bounded drain; if the message boundary
-cannot be reached, `Exchange.Wait` reports a truncated attempt.
+cannot be reached, `Exchange.Commit` reports a truncated attempt. To reject a
+capture after inspecting its response, call `Exchange.Discard`; it closes any
+unread response body and releases the captured temporary data.
 
 HTTP exchanges are written in request-then-response order by default. Set
 `rotator.UseInternetArchiveRecordOrder = true` for IA-compatible

capture_recorder.go

@@ -67,21 +67,21 @@ func (a *transportCaptureAttempt) Finish(result fhttp.CaptureAttemptResult) {
 		// A capture-sink failure cannot produce a trustworthy WARC record.
 		// Truncated network responses, however, are retained as partial records.
 		if result.Outcome == fhttp.CaptureOutcomeFailed {
-			err := errors.Join(attemptErr, a.request.close(), a.response.close())
-			a.state.finishAttempt(AttemptResult{Protocol: string(a.meta.Protocol), Outcome: result.Outcome, Err: err})
+			cleanupErr := errors.Join(a.request.close(), a.response.close())
+			a.state.finishAttempt(AttemptResult{Protocol: string(a.meta.Protocol), Outcome: result.Outcome, Err: attemptErr, cleanupErr: cleanupErr})
 			return
 		}
 
 		requestFile, err := a.request.committedFile()
 		if err != nil {
-			_ = a.response.close()
-			a.state.finishAttempt(AttemptResult{Protocol: string(a.meta.Protocol), Outcome: result.Outcome, Err: err})
+			cleanupErr := errors.Join(a.request.close(), a.response.close())
+			a.state.finishAttempt(AttemptResult{Protocol: string(a.meta.Protocol), Outcome: result.Outcome, Err: err, cleanupErr: cleanupErr})
 			return
 		}
 		responseFile, err := a.response.committedFile()
 		if err != nil {
-			_ = requestFile.Close()
-			a.state.finishAttempt(AttemptResult{Protocol: string(a.meta.Protocol), Outcome: result.Outcome, Err: err})
+			cleanupErr := errors.Join(requestFile.Close(), a.response.close())
+			a.state.finishAttempt(AttemptResult{Protocol: string(a.meta.Protocol), Outcome: result.Outcome, Err: err, cleanupErr: cleanupErr})
 			return
 		}
 
@@ -96,8 +96,13 @@ func (a *transportCaptureAttempt) Finish(result fhttp.CaptureAttemptResult) {
 			pi.Protocol = "http/1.1"
 		}
 		// Record construction and WARC I/O can outlive RoundTrip. Keep them off
-		// the transport goroutine; Exchange.Wait owns the durable completion.
+		// the transport goroutine; Exchange.Commit owns durable completion.
 		go func() {
+			if a.state.waitForDecision() == exchangeDiscard {
+				closeErr := errors.Join(requestFile.Close(), responseFile.Close())
+				a.state.finishAttempt(AttemptResult{Protocol: pi.Protocol, Outcome: result.Outcome, cleanupErr: closeErr})
+				return
+			}
 			events, writeErr := a.owner.writeCapturedExchange(a.meta.Request.Context(), a.meta.Request.URL.Scheme, requestFile, responseFile, pi, result)
 			a.state.finishAttempt(AttemptResult{Protocol: pi.Protocol, Outcome: result.Outcome, Records: events, Err: errors.Join(attemptErr, writeErr)})
 		}()

client.go

@@ -98,6 +98,7 @@ type CustomHTTPClient struct {
 	shutdownOnce             sync.Once
 	shutdownDone             chan struct{}
 	shutdownResult           FinalizeResult
+	pendingExchanges         map[*exchangeState]struct{}
 
 	CDXDedupeTotalBytes          *atomic.Int64
 	DoppelgangerDedupeTotalBytes *atomic.Int64
@@ -131,8 +132,16 @@ func (c *CustomHTTPClient) Shutdown(ctx context.Context) (FinalizeResult, error)
 	c.shutdownOnce.Do(func() {
 		c.lifecycleMu.Lock()
 		c.closing = true
+		pending := make([]*exchangeState, 0, len(c.pendingExchanges))
+		for state := range c.pendingExchanges {
+			pending = append(pending, state)
+		}
 		c.lifecycleMu.Unlock()
-		go c.runShutdown()
+		stopTransportEarly := false
+		for _, state := range pending {
+			stopTransportEarly = state.closeForShutdown() || stopTransportEarly
+		}
+		go c.runShutdown(stopTransportEarly)
 	})
 	select {
 	case <-c.shutdownDone:
@@ -144,13 +153,17 @@ func (c *CustomHTTPClient) Shutdown(ctx context.Context) (FinalizeResult, error)
 	}
 }
 
-func (c *CustomHTTPClient) runShutdown() {
+func (c *CustomHTTPClient) runShutdown(stopTransportEarly bool) {
 	defer close(c.shutdownDone)
 	var wg sync.WaitGroup
-	if c.protoClient != nil {
+	if stopTransportEarly && c.protoClient != nil {
 		c.protoClient.Shutdown()
 	}
 	c.WaitGroup.Wait()
+	c.compatWG.Wait()
+	if !stopTransportEarly && c.protoClient != nil {
+		c.protoClient.Shutdown()
+	}
 
 	close(c.WARCWriter)
 
@@ -171,7 +184,6 @@ func (c *CustomHTTPClient) runShutdown() {
 		finalizedFiles = append(finalizedFiles, result.FinalizedFiles...)
 	}
 
-	c.compatWG.Wait()
 	close(c.ErrChan)
 
 	if c.randomLocalIP {
@@ -197,14 +209,14 @@ func (c *CustomHTTPClient) Do(req *http.Request) (*http.Response, error) {
 	}
 	c.compatWG.Add(1)
 	c.lifecycleMu.Unlock()
-	exchange, err := c.Start(req)
+	exchange, err := c.start(req, exchangeCommit)
 	if exchange == nil {
 		c.compatWG.Done()
 		return nil, err
 	}
 	go func() {
 		defer c.compatWG.Done()
-		result, _ := exchange.Wait(context.Background())
+		result, _ := exchange.wait(context.Background())
 		var archiveErrs []error
 		for _, attempt := range result.Attempts {
 			archiveErrs = append(archiveErrs, attempt.Err)
@@ -214,16 +226,20 @@ func (c *CustomHTTPClient) Do(req *http.Request) (*http.Response, error) {
 			return
 		}
 		select {
-		case c.ErrChan <- &Error{Err: archiveErr, Func: "Exchange.Wait"}:
+		case c.ErrChan <- &Error{Err: archiveErr, Func: "Exchange.Commit"}:
 		default:
 		}
 	}()
 	return exchange.Response, err
 }
 
-// Start executes req and returns an Exchange whose Wait method reports the
-// durable archival result independently from receiving response headers.
+// Start executes req and returns an Exchange that the caller must finish with
+// Commit or Discard after inspecting the response.
 func (c *CustomHTTPClient) Start(req *http.Request) (*Exchange, error) {
+	return c.start(req, exchangeUndecided)
+}
+
+func (c *CustomHTTPClient) start(req *http.Request, decision exchangeDecision) (*Exchange, error) {
 	if req == nil {
 		return nil, errors.New("warc: nil request")
 	}
@@ -245,18 +261,26 @@ func (c *CustomHTTPClient) Start(req *http.Request) (*Exchange, error) {
 			return nil, errors.New("warc: feedback channel must be buffered")
 		}
 	}
-	state := newExchangeState(c, feedback)
+	state := newExchangeState(c, feedback, decision)
+	c.pendingExchanges[state] = struct{}{}
 	ctx := context.WithValue(req.Context(), exchangeContextKey{}, state)
 	req = req.Clone(ctx)
 	req.URL.Scheme = strings.ToLower(req.URL.Scheme)
 	c.WaitGroup.Add(1)
 	c.lifecycleMu.Unlock()
 	resp, err := c.protoClient.Do(ctx, req)
 	c.WaitGroup.Done()
+	state.attachResponse(resp)
 	exchange := &Exchange{Response: resp, state: state}
 	return exchange, err
 }
 
+func (c *CustomHTTPClient) unregisterExchange(state *exchangeState) {
+	c.lifecycleMu.Lock()
+	delete(c.pendingExchanges, state)
+	c.lifecycleMu.Unlock()
+}
+
 func (c *CustomHTTPClient) Get(url string) (*http.Response, error) {
 	req, err := http.NewRequest(http.MethodGet, url, nil)
 	if err != nil {
@@ -303,6 +327,7 @@ func NewWARCWritingHTTPClient(HTTPClientSettings HTTPClientSettings) (httpClient
 	}
 	httpClient = new(CustomHTTPClient)
 	httpClient.shutdownDone = make(chan struct{})
+	httpClient.pendingExchanges = make(map[*exchangeState]struct{})
 
 	httpClient.DataTotal = &DataTotal
 

client_test.go

@@ -1220,71 +1220,6 @@ func TestHTTPClientDedupeEmptyPayload(t *testing.T) {
 	}
 }
 
-func TestHTTPClientSaveChannel(t *testing.T) {
-	var (
-		rotatorSettings = defaultRotatorSettings(t)
-		errWg           sync.WaitGroup
-		err             error
-	)
-
-	// init test HTTP endpoint
-	server := newTestImageServer(t, http.StatusTooManyRequests)
-	defer server.Close()
-
-	// init the HTTP client responsible for recording HTTP(s) requests / responses
-	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 err := range httpClient.ErrChan {
-			t.Errorf("Unexpected error: %v", err)
-		}
-	}()
-
-	saveChan := make(chan bool, 1)
-
-	req, err := http.NewRequest("GET", server.URL, nil)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	req = req.WithContext(WithSaveChannel(req.Context(), saveChan))
-
-	resp, err := httpClient.Do(req)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	if resp.StatusCode == http.StatusTooManyRequests {
-		close(saveChan)
-	} else {
-		saveChan <- true
-	}
-
-	_, err = io.Copy(io.Discard, resp.Body)
-	if err != nil {
-		t.Fatalf("Unexpected error reading response body: %v", err)
-	}
-
-	httpClient.Close()
-
-	files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*")
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	for _, path := range files {
-		// note: we are actually expecting nothing here, as such, 0 for expected total. This may error if 429s aren't being filtered correctly!
-		testFileSingleHashCheck(t, path, "n/a", []string{"0"}, 0, server.URL+"/")
-	}
-}
-
 func TestHTTPClientPayloadLargerThan2MB(t *testing.T) {
 	var (
 		rotatorSettings = defaultRotatorSettings(t)

dialer.go

@@ -19,19 +19,12 @@ type contextKey string
 
 const (
 	ContextKeyFeedback contextKey = "feedback"
-	ContextKeySave     contextKey = "save"
 )
 
 func WithFeedbackChannel(ctx context.Context, feedbackChan chan FeedbackEvent) context.Context {
 	return context.WithValue(ctx, ContextKeyFeedback, feedbackChan)
 }
 
-func WithSaveChannel(ctx context.Context, ch chan bool) context.Context {
-	return context.WithValue(ctx, ContextKeySave, ch)
-}
-
-var errDiscarded = errors.New("response discarded")
-
 type dnsExchanger interface {
 	ExchangeContext(ctx context.Context, m *dns.Msg, address string) (r *dns.Msg, rtt time.Duration, err error)
 }