improve reverse proxy and serverhandling

- buffer pool for IO copy
  - flush response after read, now works with event stream
  - fixed error handling for server

internal/net/http/reverseproxy/reverse_proxy_mod.go

@@ -410,15 +410,13 @@ func (p *ReverseProxy) handler(rw http.ResponseWriter, req *http.Request) {
 
 	rw.WriteHeader(res.StatusCode)
 
-	_, err = io.Copy(rw, res.Body)
+	err = U.CopyClose(U.NewContextWriter(ctx, rw), U.NewContextReader(ctx, res.Body)) // close now, instead of defer, to populate res.Trailer
 	if err != nil {
 		if !errors.Is(err, context.Canceled) {
 			p.errorHandler(rw, req, err, true)
 		}
-		res.Body.Close()
 		return
 	}
-	res.Body.Close() // close now, instead of defer, to populate res.Trailer
 
 	if len(res.Trailer) > 0 {
 		// Force chunking if we saw a response trailer.

internal/net/http/server/error.go

@@ -8,11 +8,11 @@ import (
 	"github.com/rs/zerolog"
 )
 
-func HandleError(logger *zerolog.Logger, err error) {
+func HandleError(logger *zerolog.Logger, err error, msg string) {
 	switch {
 	case err == nil, errors.Is(err, http.ErrServerClosed), errors.Is(err, context.Canceled):
 		return
 	default:
-		logger.Fatal().Err(err).Msg("server error")
+		logger.Fatal().Err(err).Msg(msg)
 	}
 }

internal/net/http/server/server.go

@@ -99,7 +99,10 @@ func (s *Server) Start(parent task.Parent) {
 	s.startTime = time.Now()
 	if s.http != nil {
 		go func() {
-			s.handleErr(s.http.ListenAndServe())
+			err := s.http.ListenAndServe()
+			if err != nil {
+				s.handleErr(err, "failed to serve http server")
+			}
 		}()
 		s.httpStarted = true
 		s.l.Info().Str("addr", s.http.Addr).Msg("server started")
@@ -109,11 +112,11 @@ func (s *Server) Start(parent task.Parent) {
 		go func() {
 			l, err := net.Listen("tcp", s.https.Addr)
 			if err != nil {
-				s.handleErr(err)
+				s.handleErr(err, "failed to listen on port")
 				return
 			}
 			defer l.Close()
-			s.handleErr(s.https.Serve(tls.NewListener(l, s.https.TLSConfig)))
+			s.handleErr(s.https.Serve(tls.NewListener(l, s.https.TLSConfig)), "failed to serve https server")
 		}()
 		s.httpsStarted = true
 		s.l.Info().Str("addr", s.https.Addr).Msgf("server started")
@@ -131,15 +134,23 @@ func (s *Server) stop() {
 	defer cancel()
 
 	if s.http != nil && s.httpStarted {
-		s.handleErr(s.http.Shutdown(ctx))
-		s.httpStarted = false
-		s.l.Info().Str("addr", s.http.Addr).Msgf("server stopped")
+		err := s.http.Shutdown(ctx)
+		if err != nil {
+			s.handleErr(err, "failed to shutdown http server")
+		} else {
+			s.httpStarted = false
+			s.l.Info().Str("addr", s.http.Addr).Msgf("server stopped")
+		}
 	}
 
 	if s.https != nil && s.httpsStarted {
-		s.handleErr(s.https.Shutdown(ctx))
-		s.httpsStarted = false
-		s.l.Info().Str("addr", s.https.Addr).Msgf("server stopped")
+		err := s.https.Shutdown(ctx)
+		if err != nil {
+			s.handleErr(err, "failed to shutdown https server")
+		} else {
+			s.httpsStarted = false
+			s.l.Info().Str("addr", s.https.Addr).Msgf("server stopped")
+		}
 	}
 }
 
@@ -147,6 +158,6 @@ func (s *Server) Uptime() time.Duration {
 	return time.Since(s.startTime)
 }
 
-func (s *Server) handleErr(err error) {
-	HandleError(&s.l, err)
+func (s *Server) handleErr(err error, msg string) {
+	HandleError(&s.l, err, msg)
 }

internal/route/rules/cache.go

@@ -41,9 +41,7 @@ func NewCache() Cache {
 
 // Release clear the contents of the Cached and returns it to the pool.
 func (c Cache) Release() {
-	for _, k := range cacheKeys {
-		delete(c, k)
-	}
+	clear(c)
 	cachePool.Put(c)
 }
 

internal/utils/io.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"io"
+	"net/http"
 	"sync"
 	"syscall"
 
@@ -37,6 +38,14 @@ type (
 	}
 )
 
+func NewContextReader(ctx context.Context, r io.Reader) *ContextReader {
+	return &ContextReader{ctx: ctx, Reader: r}
+}
+
+func NewContextWriter(ctx context.Context, w io.Writer) *ContextWriter {
+	return &ContextWriter{ctx: ctx, Writer: w}
+}
+
 func (r *ContextReader) Read(p []byte) (int, error) {
 	select {
 	case <-r.ctx.Done():
@@ -63,7 +72,7 @@ func NewPipe(ctx context.Context, r io.ReadCloser, w io.WriteCloser) *Pipe {
 }
 
 func (p *Pipe) Start() (err error) {
-	err = Copy(&p.w, &p.r)
+	err = CopyClose(&p.w, &p.r)
 	switch {
 	case
 		// NOTE: ignoring broken pipe and connection reset by peer
@@ -97,20 +106,78 @@ func (p BidirectionalPipe) Start() E.Error {
 	return b.Error()
 }
 
+var copyBufPool = sync.Pool{
+	New: func() any {
+		return make([]byte, copyBufSize)
+	},
+}
+
+type httpFlusher interface {
+	Flush() error
+}
+
+func getHttpFlusher(dst io.Writer) httpFlusher {
+	if rw, ok := dst.(http.ResponseWriter); ok {
+		return http.NewResponseController(rw)
+	}
+	return nil
+}
+
+const (
+	copyBufSize = 32 * 1024
+)
+
 // Copyright 2009 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
-// This is a copy of io.Copy with context handling
+// This is a copy of io.Copy with context and HTTP flusher handling
 // Author: yusing <yusing@6uo.me>.
-func Copy(dst *ContextWriter, src *ContextReader) (err error) {
-	size := 32 * 1024
-	if l, ok := src.Reader.(*io.LimitedReader); ok && int64(size) > l.N {
-		if l.N < 1 {
-			size = 1
+func CopyClose(dst *ContextWriter, src *ContextReader) (err error) {
+	var buf []byte
+	if l, ok := src.Reader.(*io.LimitedReader); ok {
+		size := copyBufSize
+		if int64(size) > l.N {
+			if l.N < 1 {
+				size = 1
+			} else {
+				size = int(l.N)
+			}
+		}
+		buf = make([]byte, size)
+	} else {
+		buf = copyBufPool.Get().([]byte)
+		defer copyBufPool.Put(buf)
+	}
+	// close both as soon as one of them is done
+	wCloser, wCanClose := dst.Writer.(io.Closer)
+	rCloser, rCanClose := src.Reader.(io.Closer)
+	if wCanClose || rCanClose {
+		if src.ctx == dst.ctx {
+			go func() {
+				<-src.ctx.Done()
+				if wCanClose {
+					wCloser.Close()
+				}
+				if rCanClose {
+					rCloser.Close()
+				}
+			}()
 		} else {
-			size = int(l.N)
+			if wCloser != nil {
+				go func() {
+					<-src.ctx.Done()
+					wCloser.Close()
+				}()
+			}
+			if rCloser != nil {
+				go func() {
+					<-dst.ctx.Done()
+					rCloser.Close()
+				}()
+			}
 		}
 	}
-	buf := make([]byte, size)
+	flusher := getHttpFlusher(dst.Writer)
+	canFlush := flusher != nil
 	for {
 		select {
 		case <-src.ctx.Done():
@@ -135,6 +202,16 @@ func Copy(dst *ContextWriter, src *ContextReader) (err error) {
 					err = io.ErrShortWrite
 					return
 				}
+				if canFlush {
+					err = flusher.Flush()
+					if err != nil {
+						if errors.Is(err, http.ErrNotSupported) {
+							canFlush = false
+						} else {
+							return err
+						}
+					}
+				}
 			}
 			if er != nil {
 				if er != io.EOF {
@@ -145,7 +222,3 @@ func Copy(dst *ContextWriter, src *ContextReader) (err error) {
 		}
 	}
 }
-
-func Copy2(ctx context.Context, dst io.Writer, src io.Reader) error {
-	return Copy(&ContextWriter{ctx: ctx, Writer: dst}, &ContextReader{ctx: ctx, Reader: src})
-}