From 697d318a6ba71205a9b232af68e5038d39172d2d Mon Sep 17 00:00:00 2001
From: Evgeniy Kulikov <kim@nspcc.ru>
Date: Fri, 27 Nov 2020 15:28:27 +0300
Subject: [PATCH] WriteHeader should be called once

Signed-off-by: Evgeniy Kulikov <kim@nspcc.ru>
---
 api/metrics/api.go | 12 +++++-------
 api/router.go      | 21 ++++++++++++++-------
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/api/metrics/api.go b/api/metrics/api.go
index 8a29f37c..d72297a6 100644
--- a/api/metrics/api.go
+++ b/api/metrics/api.go
@@ -41,11 +41,11 @@ type (
 	}
 
 	responseWrapper struct {
+		sync.Once
 		http.ResponseWriter
 
-		statusCode  int
-		headWritten bool
-		startTime   time.Time
+		statusCode int
+		startTime  time.Time
 	}
 )
 
@@ -202,12 +202,10 @@ func (st *HTTPStats) updateStats(api string, w http.ResponseWriter, r *http.Requ
 
 // WriteHeader - writes http status code
 func (w *responseWrapper) WriteHeader(code int) {
-	if !w.headWritten {
+	w.Do(func() {
 		w.statusCode = code
-		w.headWritten = true
-
 		w.ResponseWriter.WriteHeader(code)
-	}
+	})
 }
 
 // Flush - Calls the underlying Flush.
diff --git a/api/router.go b/api/router.go
index af42ce04..c26df71d 100644
--- a/api/router.go
+++ b/api/router.go
@@ -3,6 +3,7 @@ package api
 import (
 	"context"
 	"net/http"
+	"sync"
 
 	"github.com/google/uuid"
 	"github.com/gorilla/mux"
@@ -81,7 +82,9 @@ type (
 	mimeType string
 
 	logResponseWriter struct {
+		sync.Once
 		http.ResponseWriter
+
 		statusCode int
 	}
 )
@@ -100,8 +103,10 @@ const (
 var _ = logErrorResponse
 
 func (lrw *logResponseWriter) WriteHeader(code int) {
-	lrw.statusCode = code
-	lrw.ResponseWriter.WriteHeader(code)
+	lrw.Do(func() {
+		lrw.statusCode = code
+		lrw.ResponseWriter.WriteHeader(code)
+	})
 }
 
 func setRequestID(h http.Handler) http.Handler {
@@ -139,12 +144,14 @@ func logErrorResponse(l *zap.Logger) mux.MiddlewareFunc {
 					zap.Int("status", lw.statusCode),
 					zap.String("method", mux.CurrentRoute(r).GetName()),
 					zap.String("description", http.StatusText(lw.statusCode)))
-			} else {
-				l.Info("call method",
-					zap.Int("status", lw.statusCode),
-					zap.String("method", mux.CurrentRoute(r).GetName()),
-					zap.String("description", http.StatusText(lw.statusCode)))
+
+				return
 			}
+
+			l.Info("call method",
+				zap.Int("status", lw.statusCode),
+				zap.String("method", mux.CurrentRoute(r).GetName()),
+				zap.String("description", http.StatusText(lw.statusCode)))
 		})
 	}
 }