chore: refactor httpx.TimeoutHandler (#3400)

rest/handler/timeouthandler.go

@@ -127,18 +127,29 @@ type timeoutWriter struct {
 
 var _ http.Pusher = (*timeoutWriter)(nil)
 
+// Flush implements the Flusher interface.
 func (tw *timeoutWriter) Flush() {
-	dst := tw.w.Header()
-	for k, vv := range tw.h {
-		dst[k] = vv
+	flusher, ok := tw.w.(http.Flusher)
+	if !ok {
+		return
 	}
-	if flusher, ok := tw.w.(http.Flusher); ok {
-		tw.w.Write(tw.wbuf.Bytes())
-		tw.wbuf.Reset()
-		flusher.Flush()
+
+	header := tw.w.Header()
+	for k, v := range tw.h {
+		header[k] = v
 	}
+
+	tw.w.Write(tw.wbuf.Bytes())
+	tw.wbuf.Reset()
+	flusher.Flush()
 }
 
+// Header returns the underline temporary http.Header.
+func (tw *timeoutWriter) Header() http.Header {
+	return tw.h
+}
+
+// Hijack implements the Hijacker interface.
 func (tw *timeoutWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	if hijacked, ok := tw.w.(http.Hijacker); ok {
 		return hijacked.Hijack()
@@ -147,14 +158,12 @@ func (tw *timeoutWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	return nil, nil, errors.New("server doesn't support hijacking")
 }
 
-// Header returns the underline temporary http.Header.
-func (tw *timeoutWriter) Header() http.Header { return tw.h }
-
 // Push implements the Pusher interface.
 func (tw *timeoutWriter) Push(target string, opts *http.PushOptions) error {
 	if pusher, ok := tw.w.(http.Pusher); ok {
 		return pusher.Push(target, opts)
 	}
+
 	return http.ErrNotSupported
 }
 
@@ -171,6 +180,7 @@ func (tw *timeoutWriter) Write(p []byte) (int, error) {
 	if !tw.wroteHeader {
 		tw.writeHeaderLocked(http.StatusOK)
 	}
+
 	return tw.wbuf.Write(p)
 }