wrap.go

Documentation: go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp

     1  // Copyright The OpenTelemetry Authors
     2  // SPDX-License-Identifier: Apache-2.0
     3  
     4  package otelhttp // import "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
     5  
     6  import (
     7  	"context"
     8  	"io"
     9  	"net/http"
    10  	"sync/atomic"
    11  
    12  	"go.opentelemetry.io/otel/propagation"
    13  )
    14  
    15  var _ io.ReadCloser = &bodyWrapper{}
    16  
    17  // bodyWrapper wraps a http.Request.Body (an io.ReadCloser) to track the number
    18  // of bytes read and the last error.
    19  type bodyWrapper struct {
    20  	io.ReadCloser
    21  	record func(n int64) // must not be nil
    22  
    23  	read atomic.Int64
    24  	err  error
    25  }
    26  
    27  func (w *bodyWrapper) Read(b []byte) (int, error) {
    28  	n, err := w.ReadCloser.Read(b)
    29  	n1 := int64(n)
    30  	w.read.Add(n1)
    31  	w.err = err
    32  	w.record(n1)
    33  	return n, err
    34  }
    35  
    36  func (w *bodyWrapper) Close() error {
    37  	return w.ReadCloser.Close()
    38  }
    39  
    40  var _ http.ResponseWriter = &respWriterWrapper{}
    41  
    42  // respWriterWrapper wraps a http.ResponseWriter in order to track the number of
    43  // bytes written, the last error, and to catch the first written statusCode.
    44  // TODO: The wrapped http.ResponseWriter doesn't implement any of the optional
    45  // types (http.Hijacker, http.Pusher, http.CloseNotifier, http.Flusher, etc)
    46  // that may be useful when using it in real life situations.
    47  type respWriterWrapper struct {
    48  	http.ResponseWriter
    49  	record func(n int64) // must not be nil
    50  
    51  	// used to inject the header
    52  	ctx context.Context
    53  
    54  	props propagation.TextMapPropagator
    55  
    56  	written     int64
    57  	statusCode  int
    58  	err         error
    59  	wroteHeader bool
    60  }
    61  
    62  func (w *respWriterWrapper) Header() http.Header {
    63  	return w.ResponseWriter.Header()
    64  }
    65  
    66  func (w *respWriterWrapper) Write(p []byte) (int, error) {
    67  	if !w.wroteHeader {
    68  		w.WriteHeader(http.StatusOK)
    69  	}
    70  	n, err := w.ResponseWriter.Write(p)
    71  	n1 := int64(n)
    72  	w.record(n1)
    73  	w.written += n1
    74  	w.err = err
    75  	return n, err
    76  }
    77  
    78  // WriteHeader persists initial statusCode for span attribution.
    79  // All calls to WriteHeader will be propagated to the underlying ResponseWriter
    80  // and will persist the statusCode from the first call.
    81  // Blocking consecutive calls to WriteHeader alters expected behavior and will
    82  // remove warning logs from net/http where developers will notice incorrect handler implementations.
    83  func (w *respWriterWrapper) WriteHeader(statusCode int) {
    84  	if !w.wroteHeader {
    85  		w.wroteHeader = true
    86  		w.statusCode = statusCode
    87  	}
    88  	w.ResponseWriter.WriteHeader(statusCode)
    89  }
    90
View as plain text