Skip to content

Commit

Permalink
Fix body size telemetry when ReadFrom is used (#26)
Browse files Browse the repository at this point in the history
Fix body size telemetry when ReadFrom is used.
This includes a small API break due to AccessHandler
hooks using int64 instead of int for size.
  • Loading branch information
ajlake authored Mar 26, 2020
1 parent cc970f4 commit ae2587f
Show file tree
Hide file tree
Showing 3 changed files with 173 additions and 9 deletions.
5 changes: 2 additions & 3 deletions baseapp/metrics.go
Original file line number Diff line number Diff line change
Expand Up @@ -79,9 +79,8 @@ func RegisterDefaultMetrics(registry metrics.Registry) {
})
}

// CountRequest is an hlog access handler that records metrics about the
// request.
func CountRequest(r *http.Request, status, _ int, _ time.Duration) {
// CountRequest is an AccessCallback that records metrics about the request.
func CountRequest(r *http.Request, status int, _ int64, _ time.Duration) {
registry := MetricsCtx(r.Context())

if c := registry.Get(MetricsKeyRequests); c != nil {
Expand Down
26 changes: 20 additions & 6 deletions baseapp/middleware.go
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ func DefaultMiddleware(logger zerolog.Logger, registry metrics.Registry) []func(
hlog.NewHandler(logger),
NewMetricsHandler(registry),
hlog.RequestIDHandler("rid", "X-Request-ID"),
hlog.AccessHandler(RecordRequest),
AccessHandler(RecordRequest),
hatpear.Catch(HandleRouteError),
hatpear.Recover(),
}
Expand All @@ -57,22 +57,36 @@ func NewMetricsHandler(registry metrics.Registry) func(http.Handler) http.Handle
}
}

// LogRequest is an hlog access handler that logs request information.
func LogRequest(r *http.Request, status, size int, elapsed time.Duration) {
// LogRequest is an AccessCallback that logs request information.
func LogRequest(r *http.Request, status int, size int64, elapsed time.Duration) {
hlog.FromRequest(r).Info().
Str("method", r.Method).
Str("path", r.URL.String()).
Str("client_ip", r.RemoteAddr).
Int("status", status).
Int("size", size).
Int64("size", size).
Dur("elapsed", elapsed).
Str("user_agent", r.UserAgent()).
Msg("http_request")
}

// RecordRequest is an hlog access handler that logs request information and
// RecordRequest is an AccessCallback that logs request information and
// records request metrics.
func RecordRequest(r *http.Request, status, size int, elapsed time.Duration) {
func RecordRequest(r *http.Request, status int, size int64, elapsed time.Duration) {
LogRequest(r, status, size, elapsed)
CountRequest(r, status, size, elapsed)
}

type AccessCallback func(r *http.Request, status int, size int64, duration time.Duration)

// AccessHandler returns a handler that call f after each request.
func AccessHandler(f AccessCallback) func(next http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
start := time.Now()
wrapped := WrapWriter(w)
next.ServeHTTP(wrapped, r)
f(r, wrapped.Status(), wrapped.BytesWritten(), time.Since(start))
})
}
}
151 changes: 151 additions & 0 deletions baseapp/recording_writer.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,151 @@
// Copyright 2020 Palantir Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Code sourced from the following and modified in a few ways:
// https://github.com/zenazn/goji/blob/a16712d37ba72246f71f9c8012974d46f8e61d16/web/mutil/writer_proxy.go

// Copyright (c) 2014, 2015, 2016 Carl Jackson ([email protected])
//
// MIT License
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

package baseapp

import (
"bufio"
"io"
"net"
"net/http"
)

// RecordingResponseWriter is a proxy for an http.ResponseWriter that
// counts bytes written and http status send to the underlying ResponseWriter.
type RecordingResponseWriter interface {
http.ResponseWriter

// Status returns the HTTP status of the request, or 0 if one has not
// yet been sent.
Status() int

// BytesWritten returns the total number of bytes sent to the client.
BytesWritten() int64
}

func WrapWriter(w http.ResponseWriter) RecordingResponseWriter {
_, cn := w.(http.CloseNotifier)
_, fl := w.(http.Flusher)
_, hj := w.(http.Hijacker)
_, rf := w.(io.ReaderFrom)

bp := basicRecorder{ResponseWriter: w}
if cn && fl && hj && rf {
return &fancyRecorder{bp}
}
if fl {
return &flushRecorder{bp}
}
return &bp
}

type basicRecorder struct {
http.ResponseWriter
code int
bytesWritten int64
}

func (b *basicRecorder) WriteHeader(code int) {
if b.code == 0 {
b.code = code
}
b.ResponseWriter.WriteHeader(code)
}

func (b *basicRecorder) Write(buf []byte) (int, error) {
if b.code == 0 {
b.code = http.StatusOK
}
n, err := b.ResponseWriter.Write(buf)
b.bytesWritten += int64(n)
return n, err
}

func (b *basicRecorder) Status() int {
return b.code
}

func (b *basicRecorder) BytesWritten() int64 {
return b.bytesWritten
}

// fancyRecorder is a writer that additionally satisfies http.CloseNotifier,
// http.Flusher, http.Hijacker, and io.ReaderFrom. It exists for the common case
// of wrapping the http.ResponseWriter that package http gives you, in order to
// make the proxied object support the full method set of the proxied object.
type fancyRecorder struct {
basicRecorder
}

func (f *fancyRecorder) CloseNotify() <-chan bool {
cn := f.basicRecorder.ResponseWriter.(http.CloseNotifier)
return cn.CloseNotify()
}
func (f *fancyRecorder) Flush() {
fl := f.basicRecorder.ResponseWriter.(http.Flusher)
fl.Flush()
}
func (f *fancyRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
hj := f.basicRecorder.ResponseWriter.(http.Hijacker)
return hj.Hijack()
}
func (f *fancyRecorder) ReadFrom(r io.Reader) (int64, error) {
if f.code == 0 {
f.code = http.StatusOK
}
rf := f.basicRecorder.ResponseWriter.(io.ReaderFrom)
n, err := rf.ReadFrom(r)
f.bytesWritten += n
return n, err
}

var _ http.CloseNotifier = &fancyRecorder{}
var _ http.Flusher = &fancyRecorder{}
var _ http.Hijacker = &fancyRecorder{}
var _ io.ReaderFrom = &fancyRecorder{}

type flushRecorder struct {
basicRecorder
}

func (f *flushRecorder) Flush() {
fl := f.basicRecorder.ResponseWriter.(http.Flusher)
fl.Flush()
}

var _ http.Flusher = &flushRecorder{}

0 comments on commit ae2587f

Please sign in to comment.