Remove muxConnectionManager and replace with MuxManager (#149)

* Remove muxConnectionManager and replace with MuxManager

* Fix unit tests, improve muxManager lifecycle

* Fix shutdownch fake in mux_manager_test

* Add missing channels to test

* Clean up signaling logic, remove generics

* Remove redundant IsShutdown check

* Add metrics and logging for muxmanager

* Ensure metricLabels is set, metrics are configured

* Fix mux_manager_test

* Refine connection metrics

* have mux observer ping the session to get latency stats

* Add more metrics

* More metrics, wrap yamux logger

* Move waiting for client metric close to include error case

* Change to 30s stream timeout

* Enable server log debugging

* more logs

* more logs

* Include details about the replacement connection

* Ensure grpc server restarts on error

* fmt

* Update metrics/prometheus_defs.go

Co-authored-by: Paul Glass <pnglass@gmail.com>

* Requested PR fixes

* fmt

---------

Co-authored-by: Paul Glass <pnglass@gmail.com>

Author: temporal-nick

metrics/prometheus_defs.go

@@ -11,12 +11,13 @@ var (
 
 	// /proxy/adminservice.go
 
-	AdminServiceStreamsActive      = DefaultGaugeVec("admin_service_streams_active", "Number of admin service streams open", "direction")
-	AdminServiceStreamDuration     = DefaultHistogramVec("admin_service_stream_duration", "The length of time each stream was open", "direction")
-	AdminServiceStreamsOpenedCount = DefaultCounterVec("admin_service_streams_opened_count", "Number of streams opened", "direction")
-	AdminServiceStreamsClosedCount = DefaultCounterVec("admin_service_streams_closed_count", "Number of streams closed", "direction")
-	AdminServiceStreamReqCount     = DefaultCounterVec("admin_service_stream_request_count", "Number of messages received", "direction")
-	AdminServiceStreamRespCount    = DefaultCounterVec("admin_service_stream_response_count", "Number of messages received", "direction")
+	AdminServiceStreamsActive        = DefaultGaugeVec("admin_service_streams_active", "Number of admin service streams open", "direction")
+	AdminServiceStreamDuration       = DefaultHistogramVec("admin_service_stream_duration", "The length of time each stream was open", "direction")
+	AdminServiceWaitingForConnection = DefaultGaugeVec("admin_service_waiting_for_connection", "Indicates the number of requests waiting on a client", "direction")
+	AdminServiceStreamsOpenedCount   = DefaultCounterVec("admin_service_streams_opened_count", "Number of streams opened", "direction")
+	AdminServiceStreamsClosedCount   = DefaultCounterVec("admin_service_streams_closed_count", "Number of streams closed", "direction")
+	AdminServiceStreamReqCount       = DefaultCounterVec("admin_service_stream_request_count", "Number of messages received", "direction")
+	AdminServiceStreamRespCount      = DefaultCounterVec("admin_service_stream_response_count", "Number of messages received", "direction")
 	// AdminServiceStreamTerminatedCount's labels are direction (inbound/outbound) and terminated_by (source/target)
 	AdminServiceStreamTerminatedCount = DefaultCounterVec("admin_service_stream_terminated_count", "Stream was terminated by remote server", "direction", "terminated_by")
 
@@ -35,14 +36,19 @@ var (
 	ProxyServiceStopped   = DefaultCounterVec("proxy_service_stopped", "Emitted on service shutdown", "direction")
 	ProxyServiceRestarted = DefaultCounterVec("proxy_service_restarted", "Emitted on service shutdown", "direction")
 
+	// /proxy/temporal_api_server.go
+
+	GRPCServerStarted = DefaultCounterVec("grpc_server_started", "Emits when the grpc server is started", "service_name")
+	GRPCServerStopped = DefaultCounterVec("grpc_server_stopped", "Emits when the grpc server is stopped", "service_name", "error")
+
 	// /transport/grpc.go
 	// Gratuitous hack: Until https://github.com/grpc-ecosystem/go-grpc-middleware/issues/783 is addressed,
 	// we need to register a dependent registry with constant labels applied.
 
 	GRPCOutboundClientMetrics = GetStandardGRPCClientInterceptor("outbound")
 	GRPCInboundClientMetrics  = GetStandardGRPCClientInterceptor("inbound")
 
-	// /transport/mux_connection_manager.go
+	// Mux Session
 
 	// Every yamux session has these available, so let's use them in the prometheus tags so we can clearly see each connection
 	muxSessionLabels = []string{"local_addr", "remote_addr", "mode", "config_name"}
@@ -52,9 +58,24 @@ var (
 		muxSessionLabels...)
 	MuxObserverReportCount = DefaultCounterVec("mux_observer_report_count", "Number of observer executions",
 		muxSessionLabels...)
+	MuxSessionPingError = DefaultCounterVec("mux_observer_session_ping_error", "Failed ping count",
+		muxSessionLabels...)
+	MuxSessionPingLatency = DefaultCounterVec("mux_observer_session_ping_latency", "Ping latency for the active session",
+		muxSessionLabels...)
+	MuxSessionPingSuccess = DefaultCounterVec("mux_observer_session_ping_success", "Ping successes for the active session",
+		muxSessionLabels...)
+
+	// Mux Manager
+
 	muxManagerLabels       = []string{"addr", "mode", "config_name"}
 	MuxErrors              = DefaultCounterVec("mux_errors", "Number of errors observed from mux", muxManagerLabels...)
 	MuxConnectionEstablish = DefaultCounterVec("mux_connection_establish", "Number of times mux has established", muxManagerLabels...)
+	MuxWaitingConnections  = DefaultGaugeVec("mux_waiting_connections", "Number of goroutines waiting for a mux connection", muxManagerLabels...)
+	MuxConnectionProvided  = DefaultCounterVec("mux_connection_provided", "Number of times a connection was provided from WithConnection", muxManagerLabels...)
+	MuxDialFailed          = DefaultCounterVec("mux_dial_failed", "Mux failed when dialing", muxManagerLabels...)
+	MuxDialSuccess         = DefaultCounterVec("mux_dial_success", "Mux succeeded on dial", muxManagerLabels...)
+
+	// Translation interceptor
 
 	translationLabels  = []string{"kind", "message_type"}
 	TranslationCount   = DefaultCounterVec("translation_success", "Count of message translations", translationLabels...)
@@ -74,6 +95,7 @@ func init() {
 	prometheus.MustRegister(collectors.NewGoCollector(collectors.WithGoCollectorRuntimeMetrics(collectors.MetricsAll),
 		collectors.WithoutGoCollectorRuntimeMetrics(collectors.MetricsDebug.Matcher)))
 	prometheus.MustRegister(AdminServiceStreamsActive)
+	prometheus.MustRegister(AdminServiceWaitingForConnection)
 	prometheus.MustRegister(AdminServiceStreamDuration)
 	prometheus.MustRegister(AdminServiceStreamsOpenedCount)
 	prometheus.MustRegister(AdminServiceStreamsClosedCount)
@@ -91,15 +113,27 @@ func init() {
 	prometheus.MustRegister(ProxyServiceCreated)
 	prometheus.MustRegister(ProxyServiceStopped)
 	prometheus.MustRegister(ProxyServiceRestarted)
+	prometheus.MustRegister(GRPCServerStarted)
+	prometheus.MustRegister(GRPCServerStopped)
 
 	prometheus.MustRegister(GRPCOutboundClientMetrics)
 	prometheus.MustRegister(GRPCInboundClientMetrics)
 
+	// Mux Session
 	prometheus.MustRegister(MuxSessionOpen)
 	prometheus.MustRegister(MuxStreamsActive)
 	prometheus.MustRegister(MuxObserverReportCount)
+	prometheus.MustRegister(MuxSessionPingError)
+	prometheus.MustRegister(MuxSessionPingLatency)
+	prometheus.MustRegister(MuxSessionPingSuccess)
+
+	// Mux Manager
 	prometheus.MustRegister(MuxErrors)
 	prometheus.MustRegister(MuxConnectionEstablish)
+	prometheus.MustRegister(MuxWaitingConnections)
+	prometheus.MustRegister(MuxConnectionProvided)
+	prometheus.MustRegister(MuxDialFailed)
+	prometheus.MustRegister(MuxDialSuccess)
 
 	prometheus.MustRegister(TranslationCount)
 	prometheus.MustRegister(TranslationErrors)

proxy/adminservice.go

@@ -261,21 +261,26 @@ func (s *adminServiceProxyServer) StreamWorkflowReplicationMessages(
 	}
 	logger.Info("AdminStreamReplicationMessages started.")
 	streamsActiveGauge := metrics.AdminServiceStreamsActive.WithLabelValues(directionLabel)
-	streamsActiveGauge.Inc()
-	metrics.AdminServiceStreamsOpenedCount.WithLabelValues(directionLabel).Inc()
-	defer streamsActiveGauge.Dec()
-	defer logger.Info("AdminStreamReplicationMessages stopped.")
 
 	// simply forwarding target metadata
 	outgoingContext := metadata.NewOutgoingContext(targetStreamServer.Context(), targetMetadata)
 	outgoingContext, cancel := context.WithCancel(outgoingContext)
 	defer cancel()
 
+	// The underlying adminClient will try to grab a connection when we call StreamWorkflowReplicationMessages.
+	// The connection is separately managed, so we want to see how long it takes to establish that conn.
+	metrics.AdminServiceWaitingForConnection.WithLabelValues(directionLabel).Inc()
 	sourceStreamClient, err := s.adminClient.StreamWorkflowReplicationMessages(outgoingContext)
+	metrics.AdminServiceWaitingForConnection.WithLabelValues(directionLabel).Dec()
 	if err != nil {
 		logger.Error("remoteAdminServiceClient.StreamWorkflowReplicationMessages encountered error", tag.Error(err))
 		return err
 	}
+	// We succesfully got a stream connection, so mark the stream as active
+	streamsActiveGauge.Inc()
+	metrics.AdminServiceStreamsOpenedCount.WithLabelValues(directionLabel).Inc()
+	defer streamsActiveGauge.Dec()
+	defer logger.Info("AdminStreamReplicationMessages stopped.")
 	streamStartTime := time.Now()
 
 	// When one side of the stream dies, we want to tell the other side to hang up

proxy/temporal_api_server.go

@@ -3,6 +3,7 @@ package proxy
 import (
 	"fmt"
 	"io"
+	"time"
 
 	"go.temporal.io/api/workflowservice/v1"
 	"go.temporal.io/server/api/adminservice/v1"
@@ -11,6 +12,7 @@ import (
 	"google.golang.org/grpc"
 
 	"github.com/temporalio/s2s-proxy/config"
+	"github.com/temporalio/s2s-proxy/metrics"
 	"github.com/temporalio/s2s-proxy/transport"
 )
 
@@ -36,6 +38,9 @@ func NewTemporalAPIServer(
 	logger log.Logger,
 ) *TemporalAPIServer {
 	server := grpc.NewServer(serverOptions...)
+	adminservice.RegisterAdminServiceServer(server, adminHandler)
+	workflowservice.RegisterWorkflowServiceServer(server, workflowserviceHandler)
+	metrics.GRPCServerStarted.WithLabelValues(serviceName)
 	return &TemporalAPIServer{
 		serviceName:            serviceName,
 		serverConfig:           serverConfig,
@@ -48,21 +53,25 @@ func NewTemporalAPIServer(
 }
 
 func (s *TemporalAPIServer) Start() {
-	adminservice.RegisterAdminServiceServer(s.server, s.adminHandler)
-	workflowservice.RegisterWorkflowServiceServer(s.server, s.workflowserviceHandler)
-
 	go func() {
-		if err := s.serverTransport.Serve(s.server); err != nil {
+		for !s.serverTransport.IsClosed() {
+			metrics.GRPCServerStarted.WithLabelValues(s.serviceName).Inc()
+			err := s.serverTransport.Serve(s.server)
 			if err == io.EOF {
 				// grpc server can get EOF error if grpc server relies on client side of
 				// mux connection. Given a mux connection from node A (mux client) to node B (mux server),
 				// and start a grpc server on A using mux client. If node B (mux server) closed the connection,
 				// grpc server on A can get an EOF client connection error from underlying mux connection.
 				// It should not happen if grpc server is based on mux server or normal TCP connection.
-				s.logger.Warn("grpc server received EOF error")
-			} else {
+				s.logger.Info("grpc server received EOF! Connection is closing")
+				metrics.GRPCServerStopped.WithLabelValues(s.serviceName, "eof").Inc()
+			} else if err != nil {
 				s.logger.Error("grpc server fatal error ", tag.Error(err))
+				metrics.GRPCServerStopped.WithLabelValues(s.serviceName, "unknown").Inc()
+			} else {
+				metrics.GRPCServerStopped.WithLabelValues(s.serviceName, "none").Inc()
 			}
+			time.Sleep(1 * time.Second)
 		}
 	}()
 }

transport/grpc.go transport/grpcutil/grpc.gotransport/grpc.go renamed to transport/grpcutil/grpc.go

@@ -1,4 +1,4 @@
-package transport
+package grpcutil
 
 import (
 	"context"
@@ -36,7 +36,7 @@ const (
 // The hostName syntax is defined in
 // https://github.com/grpc/grpc/blob/master/doc/naming.md.
 // e.g. to use dns resolver, a "dns:///" prefix should be applied to the target.
-func dial(hostName string, tlsConfig *tls.Config, clientMetrics *grpcprom.ClientMetrics, dialer func(ctx context.Context, addr string) (net.Conn, error)) (*grpc.ClientConn, error) {
+func Dial(hostName string, tlsConfig *tls.Config, clientMetrics *grpcprom.ClientMetrics, dialer func(ctx context.Context, addr string) (net.Conn, error)) (*grpc.ClientConn, error) {
 	var grpcSecureOpt grpc.DialOption
 	if tlsConfig == nil {
 		grpcSecureOpt = grpc.WithTransportCredentials(insecure.NewCredentials())

transport/mux/establisher.go

@@ -3,7 +3,6 @@ package mux
 import (
 	"crypto/tls"
 	"errors"
-	"math/rand/v2"
 	"net"
 	"time"
 
@@ -35,14 +34,23 @@ type shutdownCheck interface {
 	IsShutdown() bool
 }
 
-var retryPolicy = backoff.NewExponentialRetryPolicy(time.Second).
-	WithBackoffCoefficient(1.5).
-	WithMaximumInterval(30 * time.Second)
+var (
+	retryPolicy = backoff.NewExponentialRetryPolicy(time.Second).
+			WithBackoffCoefficient(1.5).
+			WithMaximumInterval(30 * time.Second)
+
+	// The establisher provider never has cleanup work, so we provide the same closed channel on CloseCh()
+	alwaysClosedCh = make(chan struct{})
+)
+
+func init() {
+	close(alwaysClosedCh)
+}
 
 // NewMuxEstablisherProvider makes an outbound call using the provided TCP settings. This constructor handles unpacking
 // the TLS config, configures the connection provider with retry and exponential backoff, and sets a disconnect
 // sleep time of 1-2 seconds.
-func NewMuxEstablisherProvider(name string, transportFn SetTransportCallback, setting config.TCPClientSetting, metricLabels []string, logger log.Logger, shutDown channel.ShutdownOnce) (*MuxProvider, error) {
+func NewMuxEstablisherProvider(name string, transportFn SetTransportCallback, setting config.TCPClientSetting, metricLabels []string, logger log.Logger, shutDown channel.ShutdownOnce) (MuxProvider, error) {
 	tlsWrapper := func(conn net.Conn) net.Conn { return conn }
 	if tlsCfg := setting.TLS; tlsCfg.IsEnabled() {
 		tlsConfig, err := encryption.GetClientTLSConfig(tlsCfg)
@@ -62,12 +70,14 @@ func NewMuxEstablisherProvider(name string, transportFn SetTransportCallback, se
 		shutdownCheck: shutDown,
 		metricLabels:  metricLabels,
 	}
-	sessionFn := func(conn net.Conn) (*yamux.Session, error) { return yamux.Client(conn, nil) }
-	disconnectFn := func() {
-		// If the server rapidly disconnects us, we don't want to get caught in a tight loop. Sleep 1-2 seconds before retry
-		time.Sleep(time.Second + time.Duration(rand.IntN(1000))*time.Millisecond)
+	sessionFn := func(conn net.Conn) (*yamux.Session, error) {
+		cfg := yamux.DefaultConfig()
+		cfg.Logger = wrapLoggerForYamux{logger: logger}
+		cfg.LogOutput = nil
+		cfg.StreamCloseTimeout = 30 * time.Second
+		return yamux.Client(conn, cfg)
 	}
-	return NewMuxProvider(name, connPv, sessionFn, disconnectFn, transportFn, metricLabels, logger, shutDown), nil
+	return NewMuxProvider(name, connPv, sessionFn, func() {}, transportFn, metricLabels, logger, shutDown), nil
 }
 
 // NewConnection makes a TCP call to establish a connection, then returns it. Retries with backoff over 30 seconds
@@ -85,17 +95,23 @@ func (p *establishingConnProvider) NewConnection() (net.Conn, error) {
 	}
 
 	onError := func(err error) bool {
-		p.logger.Error("mux client failed to dial", tag.Error(err))
+		if !p.shutdownCheck.IsShutdown() {
+			p.logger.Info("mux client failed to dial", tag.Error(err))
+		}
 		return !p.shutdownCheck.IsShutdown()
 	}
 	if err := backoff.ThrottleRetry(dialFn, retryPolicy, onError); err != nil {
-		p.logger.Error("mux client failed to dial with retry", tag.Error(err))
-		metrics.MuxErrors.WithLabelValues(p.metricLabels...).Inc()
+		if !p.shutdownCheck.IsShutdown() {
+			p.logger.Error("mux client failed to dial with retry", tag.Error(err))
+			metrics.MuxErrors.WithLabelValues(p.metricLabels...).Inc()
+		}
 		return nil, err
 	}
+	metrics.MuxConnectionEstablish.WithLabelValues(p.metricLabels...).Inc()
 	return client, nil
 }
 
-func (p *establishingConnProvider) CloseProvider() {
-	// Nothing to close on the client side, we're done.
+// CloseCh for the establisher is a no-op, because only the Conn needs to be closed
+func (p *establishingConnProvider) CloseCh() <-chan struct{} {
+	return alwaysClosedCh
 }