15370: Fix flaky test.

[arvados.git] / sdk / go / httpserver / responsewriter.go

diff --git a/sdk/go/httpserver/responsewriter.go b/sdk/go/httpserver/responsewriter.go
index e6ac4ca7a8be42d996e32c7fbc647389cc5248a4..049a3f1aae184de0ef522656741494db42764802 100644 (file)
--- a/sdk/go/httpserver/responsewriter.go
+++ b/sdk/go/httpserver/responsewriter.go
@@ -8,40 +8,81 @@ import (
        "net/http"
 )
 
-// ResponseWriter wraps http.ResponseWriter and exposes the status
+const sniffBytes = 1024
+
+type ResponseWriter interface {
+       http.ResponseWriter
+       WroteStatus() int
+       WroteBodyBytes() int
+       Sniffed() []byte
+}
+
+// responseWriter wraps http.ResponseWriter and exposes the status
 // sent, the number of bytes sent to the client, and the last write
 // error.
-type ResponseWriter struct {
+type responseWriter struct {
        http.ResponseWriter
-       wroteStatus    *int   // Last status given to WriteHeader()
-       wroteBodyBytes *int   // Bytes successfully written
-       err            *error // Last error returned from Write()
+       wroteStatus    int   // First status given to WriteHeader()
+       wroteBodyBytes int   // Bytes successfully written
+       err            error // Last error returned from Write()
+       sniffed        []byte
 }
 
 func WrapResponseWriter(orig http.ResponseWriter) ResponseWriter {
-       return ResponseWriter{orig, new(int), new(int), new(error)}
+       return &responseWriter{ResponseWriter: orig}
+}
+
+func (w *responseWriter) CloseNotify() <-chan bool {
+       if cn, ok := w.ResponseWriter.(http.CloseNotifier); ok {
+               return cn.CloseNotify()
+       }
+       return nil
 }
 
-func (w ResponseWriter) WriteHeader(s int) {
-       *w.wroteStatus = s
+func (w *responseWriter) WriteHeader(s int) {
+       if w.wroteStatus == 0 {
+               w.wroteStatus = s
+       }
+       // ...else it's too late to change the status seen by the
+       // client -- but we call the wrapped WriteHeader() anyway so
+       // it can log a warning.
        w.ResponseWriter.WriteHeader(s)
 }
 
-func (w ResponseWriter) Write(data []byte) (n int, err error) {
+func (w *responseWriter) Write(data []byte) (n int, err error) {
+       if w.wroteStatus == 0 {
+               w.WriteHeader(http.StatusOK)
+       } else if w.wroteStatus >= 400 {
+               w.sniff(data)
+       }
        n, err = w.ResponseWriter.Write(data)
-       *w.wroteBodyBytes += n
-       *w.err = err
+       w.wroteBodyBytes += n
+       w.err = err
        return
 }
 
-func (w ResponseWriter) WroteStatus() int {
-       return *w.wroteStatus
+func (w *responseWriter) WroteStatus() int {
+       return w.wroteStatus
+}
+
+func (w *responseWriter) WroteBodyBytes() int {
+       return w.wroteBodyBytes
+}
+
+func (w *responseWriter) Err() error {
+       return w.err
 }
 
-func (w ResponseWriter) WroteBodyBytes() int {
-       return *w.wroteBodyBytes
+func (w *responseWriter) sniff(data []byte) {
+       max := sniffBytes - len(w.sniffed)
+       if max <= 0 {
+               return
+       } else if max < len(data) {
+               data = data[:max]
+       }
+       w.sniffed = append(w.sniffed, data...)
 }
 
-func (w ResponseWriter) Err() error {
-       return *w.err
+func (w *responseWriter) Sniffed() []byte {
+       return w.sniffed
 }