fix(feg): Diameter closed connection recovery (#12192)

Signed-off-by: Evgeniy Makeev <evgeniym@fb.com>

feg/gateway/diameter/connection.go

@@ -18,6 +18,7 @@ import (
 	"net"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/fiorix/go-diameter/v4/diam"
 	"github.com/fiorix/go-diameter/v4/diam/avp"
@@ -33,6 +34,9 @@ type messageTypeEnum uint8
 const (
 	requestMessage messageTypeEnum = 1
 	answerMessage  messageTypeEnum = 2
+
+	connectionRecoveryInterval = time.Second
+	connectionRecoveryattempts = 6
 )
 
 // Connection is representing a diameter connection that you can
@@ -42,9 +46,12 @@ type Connection struct {
 	metadata *smpeer.Metadata
 	server   *DiameterServerConfig
 	client   *sm.Client
+	disabled bool
 	mutex    sync.Mutex
 }
 
+var disabledErr = errors.New("connection disabled")
+
 func newConnection(client *sm.Client, server *DiameterServerConfig) *Connection {
 	conn := &Connection{
 		server: server,
@@ -127,6 +134,9 @@ func (c *Connection) sendMessage(
 func (c *Connection) getDiamConnection() (diam.Conn, *smpeer.Metadata, error) {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
+	if c.disabled {
+		return c.conn, c.metadata, disabledErr
+	}
 	if c.conn != nil {
 		return c.conn, c.metadata, nil
 	}
@@ -155,16 +165,50 @@ func (c *Connection) getDiamConnection() (diam.Conn, *smpeer.Metadata, error) {
 		return nil, nil, errors.New("Could not obtain metadata from connection")
 	}
 	c.conn, c.metadata = conn, metadata
+	if cn, ok := conn.(diam.CloseNotifier); ok && cn != nil {
+		go c.connCloseNotify(cn.CloseNotify(), conn)
+	} else {
+		glog.Errorf("new diam conn (%T) %s -> %s is not CloseNotifier", conn, localAddr, c.server.Addr)
+	}
 	return conn, metadata, nil
 }
 
+func (c *Connection) connCloseNotify(cnc <-chan struct{}, conn diam.Conn) {
+	<-cnc // wait for close notifier
+	if glog.V(2) {
+		remoteAddress, localAddress := "nil", "nil"
+		if conn != nil {
+			if conn.LocalAddr() != nil {
+				localAddress = conn.LocalAddr().String()
+			}
+			if conn.LocalAddr() != nil {
+				remoteAddress = conn.RemoteAddr().String()
+			}
+		}
+		glog.Infof("notified of %s->%s connection closure", localAddress, remoteAddress)
+	}
+	if c != nil && conn != nil && c.destroyConnection(conn) {
+		// if connection was closed not by connection management functions, recover it
+		retryWaitTime := connectionRecoveryInterval
+		for i := 0; i < connectionRecoveryattempts; i++ {
+			_, _, err := c.getDiamConnection()
+			if err == nil || err == disabledErr {
+				break
+			}
+			time.Sleep(retryWaitTime)
+			retryWaitTime *= 2
+		}
+	}
+}
+
 // destroyConnection closes a bad connection. If the connection
 // passed is the same as the one stored in the locked connection, it is nullified.
 // If the passed diam connection is not the same, this probably means another go routine
 // already created a new connection - just try to close it and return.
-func (c *Connection) destroyConnection(conn diam.Conn) {
+// destroyConnection returns true if connection matches and was closed
+func (c *Connection) destroyConnection(conn diam.Conn) bool {
 	if conn == nil {
-		return
+		return false
 	}
 	c.mutex.Lock()
 	match := conn == c.conn
@@ -191,11 +235,13 @@ func (c *Connection) destroyConnection(conn diam.Conn) {
 		}
 	}
 	conn.Close()
+	return match
 }
 
 // cleanupConnection is similar to destroyConnection, but it closes & cleans up connection unconditionally
-func (c *Connection) cleanupConnection() {
+func (c *Connection) cleanupConnection(disabled bool) {
 	c.mutex.Lock()
+	c.disabled = disabled
 	conn := c.conn
 	if conn != nil {
 		c.conn = nil

feg/gateway/diameter/connection_manager.go

@@ -153,7 +153,7 @@ func (cm *ConnectionManager) cleanupConnections() {
 	glog.V(2).Info("ConnectionManager: removing all existing connections")
 	for _, c := range cm.connMap {
 		if c != nil {
-			c.cleanupConnection()
+			c.cleanupConnection(true)
 		}
 	}
 	cm.connMap = map[DiameterServerConnConfig]*Connection{}