@@ -67,6 +67,7 @@ type respLogger struct {
6767 addedInfo strings.Builder
6868 addedKeyValuePairs []interface {}
6969 startTime time.Time
70+ isTerminating bool
7071
7172 captureErrorOutput bool
7273
@@ -100,13 +101,13 @@ func DefaultStacktracePred(status int) bool {
100101const withLoggingLevel = 3
101102
102103// WithLogging wraps the handler with logging.
103- func WithLogging (handler http.Handler , pred StacktracePred ) http.Handler {
104+ func WithLogging (handler http.Handler , pred StacktracePred , isTerminatingFn func () bool ) http.Handler {
104105 return withLogging (handler , pred , func () bool {
105106 return klog .V (withLoggingLevel ).Enabled ()
106- })
107+ }, isTerminatingFn )
107108}
108109
109- func withLogging (handler http.Handler , stackTracePred StacktracePred , shouldLogRequest ShouldLogRequestPred ) http.Handler {
110+ func withLogging (handler http.Handler , stackTracePred StacktracePred , shouldLogRequest ShouldLogRequestPred , isTerminatingFn func () bool ) http.Handler {
110111 return http .HandlerFunc (func (w http.ResponseWriter , req * http.Request ) {
111112 if ! shouldLogRequest () {
112113 handler .ServeHTTP (w , req )
@@ -117,14 +118,16 @@ func withLogging(handler http.Handler, stackTracePred StacktracePred, shouldLogR
117118 if old := respLoggerFromRequest (req ); old != nil {
118119 panic ("multiple WithLogging calls!" )
119120 }
120-
121121 startTime := time .Now ()
122122 if receivedTimestamp , ok := request .ReceivedTimestampFrom (ctx ); ok {
123123 startTime = receivedTimestamp
124124 }
125125
126- rl := newLoggedWithStartTime (req , w , startTime )
127- rl .StacktraceWhen (stackTracePred )
126+ isTerminating := false
127+ if isTerminatingFn != nil {
128+ isTerminating = isTerminatingFn ()
129+ }
130+ rl := newLoggedWithStartTime (req , w , startTime ).StacktraceWhen (stackTracePred ).IsTerminating (isTerminating )
128131 req = req .WithContext (context .WithValue (ctx , respLoggerContextKey , rl ))
129132
130133 var logFunc func ()
@@ -135,6 +138,9 @@ func withLogging(handler http.Handler, stackTracePred StacktracePred, shouldLogR
135138 }
136139 }()
137140
141+ if klog .V (3 ).Enabled () || (rl .isTerminating && klog .V (1 ).Enabled ()) {
142+ defer rl .Log ()
143+ }
138144 w = responsewriter .WrapForHTTP1Or2 (rl )
139145 handler .ServeHTTP (w , req )
140146
@@ -205,6 +211,12 @@ func (rl *respLogger) StacktraceWhen(pred StacktracePred) *respLogger {
205211 return rl
206212}
207213
214+ // IsTerminating informs the logger that the server is terminating.
215+ func (rl * respLogger ) IsTerminating (is bool ) * respLogger {
216+ rl .isTerminating = is
217+ return rl
218+ }
219+
208220// StatusIsNot returns a StacktracePred which will cause stacktraces to be logged
209221// for any status *not* in the given list.
210222func StatusIsNot (statuses ... int ) StacktracePred {
0 commit comments