@@ -67,6 +67,7 @@ type respLogger struct {
67
67
addedInfo strings.Builder
68
68
addedKeyValuePairs []interface {}
69
69
startTime time.Time
70
+ isTerminating bool
70
71
71
72
captureErrorOutput bool
72
73
@@ -100,13 +101,13 @@ func DefaultStacktracePred(status int) bool {
100
101
const withLoggingLevel = 3
101
102
102
103
// 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 {
104
105
return withLogging (handler , pred , func () bool {
105
106
return klog .V (withLoggingLevel ).Enabled ()
106
- })
107
+ }, isTerminatingFn )
107
108
}
108
109
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 {
110
111
return http .HandlerFunc (func (w http.ResponseWriter , req * http.Request ) {
111
112
if ! shouldLogRequest () {
112
113
handler .ServeHTTP (w , req )
@@ -117,14 +118,16 @@ func withLogging(handler http.Handler, stackTracePred StacktracePred, shouldLogR
117
118
if old := respLoggerFromRequest (req ); old != nil {
118
119
panic ("multiple WithLogging calls!" )
119
120
}
120
-
121
121
startTime := time .Now ()
122
122
if receivedTimestamp , ok := request .ReceivedTimestampFrom (ctx ); ok {
123
123
startTime = receivedTimestamp
124
124
}
125
125
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 )
128
131
req = req .WithContext (context .WithValue (ctx , respLoggerContextKey , rl ))
129
132
130
133
var logFunc func ()
@@ -135,6 +138,9 @@ func withLogging(handler http.Handler, stackTracePred StacktracePred, shouldLogR
135
138
}
136
139
}()
137
140
141
+ if klog .V (3 ).Enabled () || (rl .isTerminating && klog .V (1 ).Enabled ()) {
142
+ defer rl .Log ()
143
+ }
138
144
w = responsewriter .WrapForHTTP1Or2 (rl )
139
145
handler .ServeHTTP (w , req )
140
146
@@ -205,6 +211,12 @@ func (rl *respLogger) StacktraceWhen(pred StacktracePred) *respLogger {
205
211
return rl
206
212
}
207
213
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
+
208
220
// StatusIsNot returns a StacktracePred which will cause stacktraces to be logged
209
221
// for any status *not* in the given list.
210
222
func StatusIsNot (statuses ... int ) StacktracePred {
0 commit comments