fix: reconnect bug (#19)

Author: nekomeowww

.gitignore

@@ -70,3 +70,6 @@ temp/
 
 # Debug
 __debug*
+
+# Local factorio dir
+.factorio/

docker-compose.yaml

@@ -0,0 +1,11 @@
+version: '3.8'
+
+services:
+  factorio:
+    image: factoriotools/factorio
+    restart: unless-stopped
+    ports:
+      - "34197:34197/udp"  # Game port
+      - "27015:27015/tcp"  # RCON port
+    volumes:
+      - ./.factorio:/factorio

internal/rcon/rcon.go

@@ -7,15 +7,15 @@ import (
 	"errors"
 	"io"
 	"net"
+	"strings"
 	"sync"
-	"time"
+	"sync/atomic"
 
 	"github.com/cenkalti/backoff/v4"
 	"github.com/gorcon/rcon"
 	"github.com/nekomeowww/factorio-rcon-api/v2/internal/configs"
 	"github.com/nekomeowww/fo"
 	"github.com/nekomeowww/xo/logger"
-	"github.com/samber/lo"
 	"go.uber.org/fx"
 	"go.uber.org/zap"
 )
@@ -38,7 +38,7 @@ type NewRCONParams struct {
 	Logger    *logger.Logger
 }
 
-//counterfeiter:generate -o fake/rcon.go --fake-name FakeRCON . RCON
+//counterfeiter:generate -o fake/rcon.go --fake-name FakeRCON . RCON//counterfeiter:generate -o fake/rcon.go --fake-name FakeRCON . RCON
 type RCON interface {
 	Close() error
 	Execute(ctx context.Context, command string) (string, error)
@@ -47,166 +47,173 @@ type RCON interface {
 	IsReady() bool
 }
 
+var _ RCON = (*RCONConn)(nil)
+
 type RCONConn struct {
 	*rcon.Conn
 
 	host     string
 	port     string
 	password string
 
-	ready          bool
-	readinessMutex sync.RWMutex
-	mutex          sync.RWMutex
-	logger         *logger.Logger
-	ctx            context.Context
-	cancel         context.CancelFunc
+	ready         atomic.Bool
+	reconnectChan chan struct{}
+	readyChan     chan struct{}
+
+	mutex  sync.RWMutex
+	logger *logger.Logger
+	ctx    context.Context
+	cancel context.CancelFunc
 }
 
 func NewRCON() func(NewRCONParams) (RCON, error) {
 	return func(params NewRCONParams) (RCON, error) {
 		connWrapper := &RCONConn{
-			Conn:     nil,
-			mutex:    sync.RWMutex{},
-			logger:   params.Logger,
-			host:     params.Config.Factorio.RCONHost,
-			port:     params.Config.Factorio.RCONPort,
-			password: params.Config.Factorio.RCONPassword,
+			Conn:          nil,
+			mutex:         sync.RWMutex{},
+			logger:        params.Logger,
+			host:          params.Config.Factorio.RCONHost,
+			port:          params.Config.Factorio.RCONPort,
+			password:      params.Config.Factorio.RCONPassword,
+			reconnectChan: make(chan struct{}, 1),
+			readyChan:     make(chan struct{}, 1),
 		}
 
 		ctx, cancel := context.WithCancel(context.Background())
 		connWrapper.ctx = ctx
 		connWrapper.cancel = cancel
 
-		go connWrapper.Connect(ctx)
-
-		params.Lifecycle.Append(fx.Hook{
-			OnStop: func(context.Context) error {
-				connWrapper.cancel()
+		// Start the connection manager
+		go connWrapper.connectionManager()
 
-				connWrapper.mutex.RLock()
-				defer connWrapper.mutex.RUnlock()
+		// Trigger initial connection
+		select {
+		case connWrapper.reconnectChan <- struct{}{}:
+		default:
+		}
 
-				if connWrapper.Conn == nil {
+		params.Lifecycle.Append(fx.Hook{
+			OnStop: func(ctx context.Context) error {
+				return fo.Invoke0(ctx, func() error {
+					connWrapper.cancel()
+					close(connWrapper.reconnectChan)
+					close(connWrapper.readyChan)
+
+					connWrapper.mutex.Lock()
+					defer connWrapper.mutex.Unlock()
+
+					if connWrapper.Conn != nil {
+						return connWrapper.Conn.Close()
+					}
 					return nil
-				}
-
-				_ = connWrapper.Conn.Close()
-
-				return nil
+				})
 			},
 		})
 
 		return connWrapper, nil
 	}
 }
 
-func (r *RCONConn) connect() (*rcon.Conn, error) {
-	conn, err := rcon.Dial(net.JoinHostPort(r.host, r.port), r.password)
-	if err != nil {
-		r.logger.Error("failed to connect to RCON, will attempt to reconnect", zap.Error(err))
-
-		return nil, err
-	}
+func (r *RCONConn) connectionManager() {
+	backoffStrategy := backoff.NewExponentialBackOff()
 
-	return conn, nil
-}
+	for {
+		select {
+		case <-r.ctx.Done():
+			return
+		case <-r.reconnectChan:
+			r.ready.Store(false)
 
-func (r *RCONConn) Connect(ctx context.Context) {
-	r.setUnready()
+			err := fo.Invoke0(r.ctx, func() error {
+				return backoff.Retry(func() error {
+					return r.establishConnection(r.ctx)
+				}, backoffStrategy)
+			})
 
-	err := fo.Invoke0(ctx, func() error {
-		return backoff.Retry(func() error {
-			conn, err := r.connect()
 			if err != nil {
-				return err
+				r.logger.Error("failed to establish RCON connection after retries", zap.Error(err))
+				continue
 			}
 
-			r.mutex.Lock()
-			defer r.mutex.Unlock()
-
-			r.Conn = conn
-
-			err = r.ping(ctx)
-			if err != nil {
-				r.logger.Error("failed to ping RCON, will attempt to reconnect", zap.Error(err))
+			r.ready.Store(true)
 
-				return err
+			select {
+			case r.readyChan <- struct{}{}:
+			default:
 			}
-
-			return nil
-		}, backoff.NewExponentialBackOff())
-	})
-	if err != nil {
-		r.logger.Error("failed to connect to RCON", zap.Error(err))
-		return
+		}
 	}
-
-	r.setReady()
 }
 
-func (r *RCONConn) Execute(ctx context.Context, command string) (string, error) {
-	return fo.Invoke(ctx, func() (string, error) {
-		_, _, err := lo.AttemptWithDelay(40, 250*time.Millisecond, func(_ int, _ time.Duration) error {
-			if !r.IsReady() {
-				return ErrTimeout
-			}
+func (r *RCONConn) establishConnection(ctx context.Context) error {
+	return fo.Invoke0(ctx, func() error {
+		r.mutex.Lock()
+		defer r.mutex.Unlock()
 
-			return nil
-		})
-		if err != nil {
-			return "", err
+		if r.Conn != nil {
+			_ = r.Conn.Close()
 		}
 
-		resp, err := r.Conn.Execute(command)
+		conn, err := rcon.Dial(net.JoinHostPort(r.host, r.port), r.password)
 		if err != nil {
-			if errors.Is(err, io.EOF) {
-				r.logger.Warn("RCON connection is closed, attempting to reconnect")
-
-				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-				defer cancel()
+			r.logger.Error("failed to connect to RCON", zap.Error(err))
+			return err
+		}
 
-				r.Connect(ctx)
-				return r.Execute(ctx, command)
-			}
+		r.Conn = conn
 
-			return "", err
+		// Test the connection
+		_, err = r.Conn.Execute("/help")
+		if err != nil {
+			r.logger.Error("failed to ping RCON", zap.Error(err))
+			return err
 		}
 
-		return resp, nil
+		r.logger.Info("RCON connection established successfully")
+
+		return nil
 	})
 }
 
-func (r *RCONConn) setUnready() {
-	r.readinessMutex.Lock()
-	defer r.readinessMutex.Unlock()
+func (r *RCONConn) Execute(ctx context.Context, command string) (string, error) {
+	return fo.Invoke(ctx, func() (string, error) {
+		if !r.IsReady() {
+			select {
+			case <-ctx.Done():
+				return "", ctx.Err()
+			case <-r.readyChan:
+			}
+		}
 
-	r.ready = false
-}
+		r.mutex.RLock()
+		conn := r.Conn
+		r.mutex.RUnlock()
+		if conn == nil {
+			return r.Execute(ctx, command)
+		}
 
-func (r *RCONConn) setReady() {
-	r.readinessMutex.Lock()
-	defer r.readinessMutex.Unlock()
+		resp, err := conn.Execute(command)
+		if err != nil {
+			if !strings.Contains(err.Error(), "use of closed network connection") &&
+				!strings.Contains(err.Error(), "connection reset by peer") &&
+				!errors.Is(err, io.EOF) {
+				return "", err
+			}
 
-	r.ready = true
-}
+			r.logger.Warn("RCON connection lost, reconnecting...")
 
-func (r *RCONConn) ping(ctx context.Context) error {
-	return fo.Invoke0(ctx, func() error {
-		_, err := r.Conn.Execute("/help")
-		if err != nil {
-			return err
+			select {
+			case r.reconnectChan <- struct{}{}:
+			default:
+			}
+
+			return r.Execute(ctx, command)
 		}
 
-		return nil
+		return resp, nil
 	})
 }
 
 func (r *RCONConn) IsReady() bool {
-	r.readinessMutex.RLock()
-	r.mutex.RLock()
-
-	defer r.mutex.RUnlock()
-	defer r.readinessMutex.RUnlock()
-
-	return r.ready && r.Conn != nil
+	return r.ready.Load()
 }