Add IsReady method to App interface and implement connection checks

pkg/app.go

@@ -27,7 +27,6 @@ import (
 	"reflect"
 	"strings"
 	"syscall"
-
 	"time"
 
 	"github.com/golang/protobuf/proto"
@@ -127,6 +126,7 @@ type Pitaya interface {
 	GetModule(name string) (interfaces.Module, error)
 
 	GetNumberOfConnectedClients() int64
+	IsReady(ctx context.Context) bool
 }
 
 // App is the base app struct
@@ -560,3 +560,41 @@ func (app *App) RegisterRPCJob(rpcJob worker.RPCJob) error {
 func (app *App) GetNumberOfConnectedClients() int64 {
 	return app.sessionPool.GetSessionCount()
 }
+
+// IsReady checks if pitaya is ready and able to serve requests
+func (app *App) IsReady(ctx context.Context) bool {
+	if !app.IsRunning() {
+		return false
+	}
+
+	// Check NATS RPC Client connection
+	if app.rpcClient != nil {
+		if natsClient, ok := app.rpcClient.(*cluster.NatsRPCClient); ok {
+			if !natsClient.IsConnected() {
+				logger.Log.Info("pitaya is not ready")
+				return false
+			}
+		}
+	}
+
+	// Check NATS RPC Server connection
+	if app.rpcServer != nil {
+		if natsServer, ok := app.rpcServer.(*cluster.NatsRPCServer); ok {
+			if !natsServer.IsConnected() {
+				logger.Log.Info("pitaya is not ready")
+				return false
+			}
+		}
+	}
+
+	// Check ETCD connection
+	if app.serviceDiscovery != nil {
+		if !app.serviceDiscovery.IsConnected(ctx) {
+			logger.Log.Info("pitaya is not ready")
+			return false
+		}
+	}
+
+	logger.Log.Info("pitaya is ready")
+	return true
+}

pkg/cluster/etcd_service_discovery.go

@@ -722,3 +722,15 @@ func (sd *etcdServiceDiscovery) isServerTypeBlacklisted(svType string) bool {
 	}
 	return false
 }
+
+func (sd *etcdServiceDiscovery) IsConnected(ctx context.Context) bool {
+	if sd.cli == nil {
+		return false
+	}
+
+	checkCtx, cancel := context.WithTimeout(ctx, sd.revokeTimeout)
+	defer cancel()
+
+	_, err := sd.cli.Get(checkCtx, "health-check", clientv3.WithLimit(1))
+	return err == nil
+}

pkg/cluster/nats_rpc_client.go

@@ -262,7 +262,6 @@ func (ns *NatsRPCClient) Init() error {
 
 // initConnection initializes or replaces the NATS connection
 func (ns *NatsRPCClient) initConnection(isReplacement bool) error {
-
 	if !isReplacement {
 		ns.running = true
 		logger.Log.Debugf("connecting to nats (client) with timeout of %s", ns.connectionTimeout)
@@ -312,3 +311,7 @@ func (ns *NatsRPCClient) stop() {
 func (ns *NatsRPCClient) getSubscribeChannel() string {
 	return fmt.Sprintf("pitaya/servers/%s/%s", ns.server.Type, ns.server.ID)
 }
+
+func (ns *NatsRPCClient) IsConnected() bool {
+	return ns.conn != nil && ns.conn.IsConnected()
+}

pkg/cluster/nats_rpc_server.go

@@ -356,7 +356,6 @@ func (ns *NatsRPCServer) Init() error {
 
 // initConnection initializes or replaces the NATS connection
 func (ns *NatsRPCServer) initConnection(isReplacement bool) error {
-
 	if !isReplacement {
 		// TODO should we have concurrency here? it feels like we should
 		go ns.handleMessages()
@@ -479,3 +478,8 @@ func (ns *NatsRPCServer) reportMetrics() {
 		}
 	}
 }
+
+// IsConnected returns true if NATS connection is established
+func (ns *NatsRPCServer) IsConnected() bool {
+	return ns.conn != nil && ns.conn.IsConnected()
+}

pkg/cluster/service_discovery.go

@@ -20,7 +20,11 @@
 
 package cluster
 
-import "github.com/topfreegames/pitaya/v3/pkg/interfaces"
+import (
+	"context"
+
+	"github.com/topfreegames/pitaya/v3/pkg/interfaces"
+)
 
 // ServiceDiscovery is the interface for a service discovery client
 type ServiceDiscovery interface {
@@ -29,5 +33,6 @@ type ServiceDiscovery interface {
 	GetServers() []*Server
 	SyncServers(firstSync bool) error
 	AddListener(listener SDListener)
+	IsConnected(ctx context.Context) bool
 	interfaces.Module
 }

pkg/static.go

@@ -229,3 +229,7 @@ func GetModule(name string) (interfaces.Module, error) {
 func GetNumberOfConnectedClients() int64 {
 	return DefaultApp.GetNumberOfConnectedClients()
 }
+
+func IsReady(ctx context.Context) bool {
+	return DefaultApp.IsReady(ctx)
+}