Add KeepAliveMaxTime and KeepAliveMaxRequests features to entrypoints

pkg/server/server_entrypoint_tcp.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"errors"
+	"expvar"
 	"fmt"
 	stdlog "log"
 	"net"
@@ -34,6 +35,25 @@ import (
 
 var httpServerLogger = stdlog.New(log.WithoutContext().WriterLevel(logrus.DebugLevel), "", 0)
 
+type key string
+
+const (
+	connStateKey       key    = "connState"
+	debugConnectionEnv string = "DEBUG_CONNECTION"
+)
+
+var (
+	clientConnectionStates   = map[string]*connState{}
+	clientConnectionStatesMu = sync.RWMutex{}
+)
+
+type connState struct {
+	State            string
+	KeepAliveState   string
+	Start            time.Time
+	HTTPRequestCount int
+}
+
 type httpForwarder struct {
 	net.Listener
 	connChan chan net.Conn
@@ -68,6 +88,12 @@ type TCPEntryPoints map[string]*TCPEntryPoint
 
 // NewTCPEntryPoints creates a new TCPEntryPoints.
 func NewTCPEntryPoints(entryPointsConfig static.EntryPoints, hostResolverConfig *types.HostResolverConfig) (TCPEntryPoints, error) {
+	if os.Getenv(debugConnectionEnv) != "" {
+		expvar.Publish("clientConnectionStates", expvar.Func(func() any {
+			return clientConnectionStates
+		}))
+	}
+
 	serverEntryPointsTCP := make(TCPEntryPoints)
 	for entryPointName, config := range entryPointsConfig {
 		protocol, err := config.GetProtocol()
@@ -548,6 +574,11 @@ func createHTTPServer(ctx context.Context, ln net.Listener, configuration *stati
 		})
 	}
 
+	debugConnection := os.Getenv(debugConnectionEnv) != ""
+	if debugConnection || (configuration.Transport != nil && (configuration.Transport.KeepAliveMaxTime > 0 || configuration.Transport.KeepAliveMaxRequests > 0)) {
+		handler = newKeepAliveMiddleware(handler, configuration.Transport.KeepAliveMaxRequests, configuration.Transport.KeepAliveMaxTime)
+	}
+
 	serverHTTP := &http.Server{
 		Handler:      handler,
 		ErrorLog:     httpServerLogger,
@@ -555,6 +586,27 @@ func createHTTPServer(ctx context.Context, ln net.Listener, configuration *stati
 		WriteTimeout: time.Duration(configuration.Transport.RespondingTimeouts.WriteTimeout),
 		IdleTimeout:  time.Duration(configuration.Transport.RespondingTimeouts.IdleTimeout),
 	}
+	if debugConnection || (configuration.Transport != nil && (configuration.Transport.KeepAliveMaxTime > 0 || configuration.Transport.KeepAliveMaxRequests > 0)) {
+		serverHTTP.ConnContext = func(ctx context.Context, c net.Conn) context.Context {
+			cState := &connState{Start: time.Now()}
+			if debugConnection {
+				clientConnectionStatesMu.Lock()
+				clientConnectionStates[getConnKey(c)] = cState
+				clientConnectionStatesMu.Unlock()
+			}
+			return context.WithValue(ctx, connStateKey, cState)
+		}
+
+		if debugConnection {
+			serverHTTP.ConnState = func(c net.Conn, state http.ConnState) {
+				clientConnectionStatesMu.Lock()
+				if clientConnectionStates[getConnKey(c)] != nil {
+					clientConnectionStates[getConnKey(c)].State = state.String()
+				}
+				clientConnectionStatesMu.Unlock()
+			}
+		}
+	}
 
 	// ConfigureServer configures HTTP/2 with the MaxConcurrentStreams option for the given server.
 	// Also keeping behavior the same as
@@ -584,6 +636,10 @@ func createHTTPServer(ctx context.Context, ln net.Listener, configuration *stati
 	}, nil
 }
 
+func getConnKey(conn net.Conn) string {
+	return fmt.Sprintf("%s => %s", conn.RemoteAddr(), conn.LocalAddr())
+}
+
 func newTrackedConnection(conn tcp.WriteCloser, tracker *connectionTracker) *trackedConnection {
 	tracker.AddConnection(conn)
 	return &trackedConnection{