chore: import upstream snapshot with attribution

internal/distributed/proxy/httpserver/timeout_middleware.go

@@ -0,0 +1,326 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package httpserver
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"fmt"
+	"net"
+	"net/http"
+	"strconv"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+
+	mhttp "github.com/milvus-io/milvus/internal/http"
+	"github.com/milvus-io/milvus/internal/json"
+	"github.com/milvus-io/milvus/pkg/v3/mlog"
+	"github.com/milvus-io/milvus/pkg/v3/util/merr"
+	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
+)
+
+// BufferPool represents a pool of buffers.
+type BufferPool struct {
+	pool sync.Pool
+}
+
+// Get returns a buffer from the buffer pool.
+// If the pool is empty, a new buffer is created and returned.
+func (p *BufferPool) Get() *bytes.Buffer {
+	buf := p.pool.Get()
+	if buf == nil {
+		return &bytes.Buffer{}
+	}
+	return buf.(*bytes.Buffer)
+}
+
+// Put adds a buffer back to the pool.
+func (p *BufferPool) Put(buf *bytes.Buffer) {
+	p.pool.Put(buf)
+}
+
+// Timeout struct
+type Timeout struct {
+	handler gin.HandlerFunc
+}
+
+const timeoutRecorderNoWritten = -1
+
+type timeoutResponseRecorder struct {
+	body        *bytes.Buffer
+	headers     http.Header
+	mu          sync.Mutex
+	closed      bool
+	status      int
+	size        int
+	closeNotify chan bool
+}
+
+func newTimeoutResponseRecorder(buf *bytes.Buffer) *timeoutResponseRecorder {
+	return &timeoutResponseRecorder{
+		body:        buf,
+		headers:     make(http.Header),
+		status:      http.StatusOK,
+		size:        timeoutRecorderNoWritten,
+		closeNotify: make(chan bool),
+	}
+}
+
+func (w *timeoutResponseRecorder) Header() http.Header {
+	return w.headers
+}
+
+func (w *timeoutResponseRecorder) Write(data []byte) (int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.closed || w.body == nil {
+		return 0, merr.WrapErrServiceInternalMsg("response writer closed")
+	}
+	if !w.written() {
+		w.size = 0
+	}
+	n, err := w.body.Write(data)
+	w.size += n
+	return n, err
+}
+
+func (w *timeoutResponseRecorder) WriteString(s string) (int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.closed || w.body == nil {
+		return 0, merr.WrapErrServiceInternalMsg("response writer closed")
+	}
+	if !w.written() {
+		w.size = 0
+	}
+	n, err := w.body.WriteString(s)
+	w.size += n
+	return n, err
+}
+
+func (w *timeoutResponseRecorder) WriteHeader(code int) {
+	if code == -1 {
+		return
+	}
+	checkWriteHeaderCode(code)
+
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.closed || w.written() {
+		return
+	}
+	w.status = code
+}
+
+func (w *timeoutResponseRecorder) WriteHeaderNow() {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.closed || w.written() {
+		return
+	}
+	w.size = 0
+}
+
+func (w *timeoutResponseRecorder) Status() int {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.status
+}
+
+func (w *timeoutResponseRecorder) Size() int {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.size
+}
+
+func (w *timeoutResponseRecorder) Written() bool {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.written()
+}
+
+func (w *timeoutResponseRecorder) Flush() {}
+
+func (w *timeoutResponseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return nil, nil, merr.WrapErrServiceInternalMsg("response writer does not support hijack")
+}
+
+func (w *timeoutResponseRecorder) CloseNotify() <-chan bool {
+	return w.closeNotify
+}
+
+func (w *timeoutResponseRecorder) Pusher() http.Pusher {
+	return nil
+}
+
+func (w *timeoutResponseRecorder) CommitTo(realWriter gin.ResponseWriter) error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.closed || w.body == nil {
+		return merr.WrapErrServiceInternalMsg("response writer closed")
+	}
+
+	dst := realWriter.Header()
+	for k, vv := range w.headers {
+		dst[k] = append([]string(nil), vv...)
+	}
+	realWriter.WriteHeader(w.status)
+	if w.body.Len() == 0 {
+		if w.written() || w.status != http.StatusOK {
+			realWriter.WriteHeaderNow()
+		}
+		return nil
+	}
+	_, err := realWriter.Write(w.body.Bytes())
+	return err
+}
+
+func (w *timeoutResponseRecorder) Close() {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.body != nil {
+		w.body.Reset()
+		w.body = nil
+	}
+	w.closed = true
+}
+
+func (w *timeoutResponseRecorder) CloseForTimeout() {
+	w.Close()
+}
+
+func (w *timeoutResponseRecorder) written() bool {
+	return w.size != timeoutRecorderNoWritten
+}
+
+func checkWriteHeaderCode(code int) {
+	if code < 100 || code > 999 {
+		panic(fmt.Sprintf("invalid http status code: %d", code))
+	}
+}
+
+var timeoutContextKeysToPropagate = []string{
+	HTTPReturnCode,
+	HTTPReturnMessage,
+	ContextRequest,
+	ContextResponse,
+	"traceID",
+}
+
+func propagateTimeoutContextKeys(dst *gin.Context, src *gin.Context) {
+	for _, key := range timeoutContextKeysToPropagate {
+		if value, ok := src.Get(key); ok {
+			dst.Set(key, value)
+		}
+	}
+}
+
+func timeoutMiddleware(handler gin.HandlerFunc) gin.HandlerFunc {
+	timeoutHandler := &Timeout{
+		handler: handler,
+	}
+	bufPool := &BufferPool{}
+	return func(gCtx *gin.Context) {
+		timeout := paramtable.Get().HTTPCfg.RequestTimeoutMs.GetAsDuration(time.Millisecond)
+		requestTimeout := gCtx.Request.Header.Get(mhttp.HTTPHeaderRequestTimeout)
+		if requestTimeout != "" {
+			timeoutSecond, err := strconv.ParseInt(requestTimeout, 10, 64)
+			if err != nil {
+				HTTPAbortReturn(gCtx, http.StatusOK, gin.H{
+					mhttp.HTTPReturnCode: merr.Code(merr.ErrParameterInvalid),
+					mhttp.HTTPReturnMessage: merr.WrapErrParameterInvalidMsg(
+						"%s parse failed, err: %s",
+						mhttp.HTTPHeaderRequestTimeout,
+						err.Error(),
+					).Error(),
+				})
+				return
+			}
+			timeout = time.Duration(timeoutSecond) * time.Second
+		}
+		topCtx, cancel := context.WithTimeout(gCtx.Request.Context(), timeout)
+		defer cancel()
+		req := gCtx.Request.WithContext(topCtx)
+		gCtx.Request = req
+
+		finish := make(chan struct{}, 1)
+		panicChan := make(chan interface{}, 1)
+
+		realWriter := gCtx.Writer
+		buffer := bufPool.Get()
+		buffer.Reset()
+		recorder := newTimeoutResponseRecorder(buffer)
+		handlerCtx := gCtx.Copy()
+		handlerCtx.Request = req
+		handlerCtx.Writer = recorder
+
+		go func() {
+			defer func() {
+				if p := recover(); p != nil {
+					panicChan <- p
+				}
+			}()
+			timeoutHandler.handler(handlerCtx)
+			finish <- struct{}{}
+		}()
+
+		timer := time.NewTimer(timeout)
+		defer timer.Stop()
+
+		select {
+		case p := <-panicChan:
+			recorder.Close()
+			bufPool.Put(buffer)
+			gCtx.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{mhttp.HTTPReturnCode: http.StatusInternalServerError})
+			panic(p)
+
+		case <-finish:
+			propagateTimeoutContextKeys(gCtx, handlerCtx)
+			if handlerCtx.IsAborted() {
+				gCtx.Abort()
+			}
+			gCtx.Next()
+			if err := recorder.CommitTo(realWriter); err != nil {
+				mlog.Warn(context.TODO(), "failed to write response body", mlog.Err(err))
+				recorder.Close()
+				bufPool.Put(buffer)
+				return
+			}
+			recorder.Close()
+			bufPool.Put(buffer)
+
+		case <-timer.C:
+			cancel()
+			gCtx.Abort()
+			gCtx.Set(HTTPReturnCode, merr.TimeoutCode)
+			gCtx.Set(HTTPReturnMessage, "request timeout")
+			recorder.CloseForTimeout()
+			bufPool.Put(buffer)
+
+			realWriter.Header().Set("Content-Type", "application/json; charset=utf-8")
+			if traceID, ok := getTraceID(gCtx); ok {
+				setTraceIDHeaderTo(realWriter.Header(), traceID)
+			}
+			realWriter.WriteHeader(http.StatusRequestTimeout)
+			body, _ := json.Marshal(gin.H{HTTPReturnCode: merr.TimeoutCode, HTTPReturnMessage: "request timeout"})
+			realWriter.Write(body)
+		}
+	}
+}