From ae2587f34a1c067f8b43f7c7b71561825e8eeaa9 Mon Sep 17 00:00:00 2001 From: ajlake Date: Thu, 26 Mar 2020 15:59:25 -0700 Subject: [PATCH] Fix body size telemetry when ReadFrom is used (#26) 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. --- baseapp/metrics.go | 5 +- baseapp/middleware.go | 26 +++++-- baseapp/recording_writer.go | 151 ++++++++++++++++++++++++++++++++++++ 3 files changed, 173 insertions(+), 9 deletions(-) create mode 100644 baseapp/recording_writer.go diff --git a/baseapp/metrics.go b/baseapp/metrics.go index a6ae8933..1dc7f34b 100644 --- a/baseapp/metrics.go +++ b/baseapp/metrics.go @@ -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 { diff --git a/baseapp/middleware.go b/baseapp/middleware.go index e2f03cb4..c622b642 100644 --- a/baseapp/middleware.go +++ b/baseapp/middleware.go @@ -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(), } @@ -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)) + }) + } +} diff --git a/baseapp/recording_writer.go b/baseapp/recording_writer.go new file mode 100644 index 00000000..91317736 --- /dev/null +++ b/baseapp/recording_writer.go @@ -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 (carl@avtok.com) +// +// 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{}