resolve merge conflict

cmd/staking-api-service/main.go

@@ -10,6 +10,7 @@ import (
 	"github.com/babylonlabs-io/staking-api-service/internal/clients"
 	"github.com/babylonlabs-io/staking-api-service/internal/config"
 	"github.com/babylonlabs-io/staking-api-service/internal/db/model"
+	"github.com/babylonlabs-io/staking-api-service/internal/observability/healthcheck"
 	"github.com/babylonlabs-io/staking-api-service/internal/observability/metrics"
 	"github.com/babylonlabs-io/staking-api-service/internal/queue"
 	"github.com/babylonlabs-io/staking-api-service/internal/services"
@@ -81,6 +82,8 @@ func main() {
 
 	queues.StartReceivingMessages()
 
+	healthcheck.StartHealthCheckCron(ctx, queues, cfg.Server.HealthCheckInterval)
+
 	apiServer, err := api.New(ctx, cfg, services)
 	if err != nil {
 		log.Fatal().Err(err).Msg("error while setting up staking api service")

config/config-docker.yml

@@ -8,6 +8,7 @@ server:
   log-level: debug
   btc-net: "mainnet"
   max-content-length: 4096
+  health-check-interval: 300 # 5 minutes interval
 db:
   username: root
   password: example

config/config-local.yml

@@ -8,6 +8,7 @@ server:
   log-level: debug
   btc-net: "signet"
   max-content-length: 4096
+  health-check-interval: 300 # 5 minutes interval
 db:
   username: root
   password: example

go.mod

@@ -13,6 +13,7 @@ require (
 	github.com/btcsuite/btcd/btcutil v1.1.5
 	github.com/btcsuite/btcd/chaincfg/chainhash v1.1.0
 	github.com/rabbitmq/amqp091-go v1.9.0
+	github.com/robfig/cron/v3 v3.0.1
 	github.com/spf13/viper v1.18.2
 	github.com/swaggo/swag v1.16.3
 	github.com/unrolled/secure v1.14.0

go.sum

@@ -1002,6 +1002,8 @@ github.com/rcrowley/go-metrics v0.0.0-20201227073835-cf1acfcdf475 h1:N/ElC8H3+5X
 github.com/rcrowley/go-metrics v0.0.0-20201227073835-cf1acfcdf475/go.mod h1:bCqnVzQkZxMG4s8nGwiZ5l3QUCyqpo9Y+/ZMZ9VjZe4=
 github.com/regen-network/protobuf v1.3.3-alpha.regen.1 h1:OHEc+q5iIAXpqiqFKeLpu5NwTIkVXUs48vFMwzqpqY4=
 github.com/regen-network/protobuf v1.3.3-alpha.regen.1/go.mod h1:2DjTFR1HhMQhiWC5sZ4OhQ3+NtdbZ6oBDKQwq5Ou+FI=
+github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs=
+github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro=
 github.com/rogpeppe/fastuuid v0.0.0-20150106093220-6724a57986af/go.mod h1:XWv6SoW27p1b0cqNHllgS5HIMJraePCO15w5zCzIWYg=
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=

internal/config/server.go

@@ -13,15 +13,16 @@ import (
 )
 
 type ServerConfig struct {
-	Host             string        `mapstructure:"host"`
-	Port             int           `mapstructure:"port"`
-	WriteTimeout     time.Duration `mapstructure:"write-timeout"`
-	ReadTimeout      time.Duration `mapstructure:"read-timeout"`
-	IdleTimeout      time.Duration `mapstructure:"idle-timeout"`
-	AllowedOrigins   []string      `mapstructure:"allowed-origins"`
-	BTCNet           string        `mapstructure:"btc-net"`
-	LogLevel         string        `mapstructure:"log-level"`
-	MaxContentLength int64         `mapstructure:"max-content-length"`
+	Host                string        `mapstructure:"host"`
+	Port                int           `mapstructure:"port"`
+	WriteTimeout        time.Duration `mapstructure:"write-timeout"`
+	ReadTimeout         time.Duration `mapstructure:"read-timeout"`
+	IdleTimeout         time.Duration `mapstructure:"idle-timeout"`
+	AllowedOrigins      []string      `mapstructure:"allowed-origins"`
+	BTCNet              string        `mapstructure:"btc-net"`
+	LogLevel            string        `mapstructure:"log-level"`
+	MaxContentLength    int64         `mapstructure:"max-content-length"`
+	HealthCheckInterval int           `mapstructure:"health-check-interval"`
 
 	BTCNetParam *chaincfg.Params
 }
@@ -52,6 +53,10 @@ func (cfg *ServerConfig) Validate() error {
 		return fmt.Errorf("MaxContentLength must be a positive integer")
 	}
 
+	if cfg.HealthCheckInterval <= 0 {
+		return fmt.Errorf("HealthCheckInterval must be a positive integer")
+	}
+
 	btcNet, err := utils.GetBtcNetParamesFromString(cfg.BTCNet)
 	if err != nil {
 		return errors.New("invalid btc-net")

internal/observability/healthcheck/healthcheck.go

@@ -0,0 +1,60 @@
+package healthcheck
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/babylonlabs-io/staking-api-service/internal/observability/metrics"
+	"github.com/babylonlabs-io/staking-api-service/internal/queue"
+	"github.com/robfig/cron/v3"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+)
+
+var logger zerolog.Logger = log.Logger
+
+func SetLogger(customLogger zerolog.Logger) {
+	logger = customLogger
+}
+
+func StartHealthCheckCron(ctx context.Context, queues *queue.Queues, cronTime int) error {
+	c := cron.New()
+	logger.Info().Msg("Initiated Health Check Cron")
+
+	if cronTime == 0 {
+		cronTime = 60
+	}
+
+	cronSpec := fmt.Sprintf("@every %ds", cronTime)
+
+	_, err := c.AddFunc(cronSpec, func() {
+		queueHealthCheck(queues)
+	})
+
+	if err != nil {
+		return err
+	}
+
+	c.Start()
+
+	go func() {
+		<-ctx.Done()
+		logger.Info().Msg("Stopping Health Check Cron")
+		c.Stop()
+	}()
+
+	return nil
+}
+
+func queueHealthCheck(queues *queue.Queues) {
+	if err := queues.IsConnectionHealthy(); err != nil {
+		logger.Error().Err(err).Msg("One or more queue connections are not healthy.")
+		// Record service unavailable in metrics
+		metrics.RecordServiceCrash("queue")
+		terminateService()
+	}
+}
+
+func terminateService() {
+	logger.Fatal().Msg("Terminating service due to health check failure.")
+}

internal/observability/metrics/metrics.go

@@ -32,6 +32,7 @@ var (
 	queueOperationFailureCounter     *prometheus.CounterVec
 	httpResponseWriteFailureCounter  *prometheus.CounterVec
 	clientRequestDurationHistogram   *prometheus.HistogramVec
+	serviceCrashCounter              *prometheus.CounterVec
 )
 
 // Init initializes the metrics package.
@@ -113,6 +114,13 @@ func registerMetrics() {
 		},
 		[]string{"baseurl", "method", "path", "status"},
 	)
+	serviceCrashCounter = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "service_crash_total",
+			Help: "",
+		},
+		[]string{"type"},
+	)
 
 	prometheus.MustRegister(
 		httpRequestDurationHistogram,
@@ -121,6 +129,7 @@ func registerMetrics() {
 		queueOperationFailureCounter,
 		httpResponseWriteFailureCounter,
 		clientRequestDurationHistogram,
+		serviceCrashCounter,
 	)
 }
 
@@ -177,3 +186,8 @@ func StartClientRequestDurationTimer(baseUrl, method, path string) func(statusCo
 		).Observe(duration)
 	}
 }
+
+// RecordServiceCrash increments the service crash counter.
+func RecordServiceCrash(service string) {
+	serviceCrashCounter.WithLabelValues(service).Inc()
+}

internal/queue/queue.go

@@ -2,7 +2,9 @@ package queue
 
 import (
 	"context"
+	"fmt"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/babylonlabs-io/staking-api-service/internal/observability/metrics"
@@ -257,3 +259,28 @@ func recordErrorLog(err *types.Error) {
 		log.Warn().Err(err).Msg("event processing failed with 4xx error")
 	}
 }
+
+func (q *Queues) IsConnectionHealthy() error {
+	var errorMessages []string
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5 * time.Second)
+	defer cancel()
+
+	checkQueue := func(name string, client client.QueueClient) {
+		if err := client.Ping(ctx); err != nil {
+			errorMessages = append(errorMessages, fmt.Sprintf("%s is not healthy: %v", name, err))
+		}
+	}
+
+	checkQueue("ActiveStakingQueueClient", q.ActiveStakingQueueClient)
+	checkQueue("ExpiredStakingQueueClient", q.ExpiredStakingQueueClient)
+	checkQueue("UnbondingStakingQueueClient", q.UnbondingStakingQueueClient)
+	checkQueue("WithdrawStakingQueueClient", q.WithdrawStakingQueueClient)
+	checkQueue("StatsQueueClient", q.StatsQueueClient)
+	checkQueue("BtcInfoQueueClient", q.BtcInfoQueueClient)
+
+	if len(errorMessages) > 0 {
+		return fmt.Errorf(strings.Join(errorMessages, "; "))
+	}
+	return nil
+}

internal/services/stats.go

@@ -17,6 +17,7 @@ type OverallStatsPublic struct {
 	TotalDelegations  int64  `json:"total_delegations"`
 	TotalStakers      uint64 `json:"total_stakers"`
 	UnconfirmedTvl    uint64 `json:"unconfirmed_tvl"`
+	PendingTvl        uint64 `json:"pending_tvl"`
 }
 
 type StakerStatsPublic struct {
@@ -131,7 +132,11 @@ func (s *Services) GetOverallStats(ctx context.Context) (*OverallStatsPublic, *t
 		log.Ctx(ctx).Error().Err(err).Msg("error while fetching overall stats")
 		return nil, types.NewInternalServiceError(err)
 	}
+
 	unconfirmedTvl := uint64(0)
+	confirmedTvl := uint64(0)
+	pendingTvl := uint64(0)
+
 	btcInfo, err := s.DbClient.GetLatestBtcInfo(ctx)
 	if err != nil {
 		// Handle missing BTC information, which may occur during initial setup.
@@ -145,15 +150,18 @@ func (s *Services) GetOverallStats(ctx context.Context) (*OverallStatsPublic, *t
 		}
 	} else {
 		unconfirmedTvl = btcInfo.UnconfirmedTvl
+		confirmedTvl = btcInfo.ConfirmedTvl
+		pendingTvl = unconfirmedTvl - confirmedTvl
 	}
 
 	return &OverallStatsPublic{
-		ActiveTvl:         stats.ActiveTvl,
+		ActiveTvl:         int64(confirmedTvl),
 		TotalTvl:          stats.TotalTvl,
 		ActiveDelegations: stats.ActiveDelegations,
 		TotalDelegations:  stats.TotalDelegations,
 		TotalStakers:      stats.TotalStakers,
 		UnconfirmedTvl:    unconfirmedTvl,
+		PendingTvl:        pendingTvl,
 	}, nil
 }
 

tests/config/config-test.yml

@@ -8,6 +8,7 @@ server:
   log-level: error
   btc-net: "signet"
   max-content-length: 40960
+  health-check-interval: 2
 db:
   username: root
   password: example

tests/healthcheck_test.go

@@ -124,4 +124,4 @@ func TestSecurityHeaders(t *testing.T) {
 	assert.Equal(t, "DENY", resp.Header.Get("X-Frame-Options"), "expected X-Frame-Options to be DENY")
 	assert.Equal(t, "default-src 'self'; script-src 'self' 'unsafe-inline' https://cdnjs.cloudflare.com https://stackpath.bootstrap.com; style-src 'self' 'unsafe-inline' https://cdnjs.cloudflare.com https://stackpath.bootstrap.com; img-src 'self' data: https://cdnjs.cloudflare.com https://stackpath.bootstrap.com; font-src 'self' https://cdnjs.cloudflare.com https://stackpath.bootstrap.com; object-src 'none'; frame-ancestors 'self'; form-action 'self'; block-all-mixed-content; base-uri 'self';", resp.Header.Get("Content-Security-Policy"), "expected Swagger Content-Security-Policy")
 	assert.Equal(t, "strict-origin-when-cross-origin", resp.Header.Get("Referrer-Policy"), "expected Referrer-Policy to be strict-origin-when-cross-origin")
-}
+}

tests/stats_test.go

@@ -204,13 +204,14 @@ func TestStatsEndpoints(t *testing.T) {
 
 	// Test the overall stats endpoint
 	overallStats := fetchOverallStatsEndpoint(t, testServer)
-	assert.Equal(t, int64(activeStakingEvent.StakingValue), overallStats.ActiveTvl)
 	assert.Equal(t, int64(activeStakingEvent.StakingValue), overallStats.TotalTvl)
 	assert.Equal(t, int64(1), overallStats.ActiveDelegations)
 	assert.Equal(t, int64(1), overallStats.TotalDelegations)
 	assert.Equal(t, uint64(1), overallStats.TotalStakers)
-	// We have not yet sent any UnconfirmedInfoEvent, hence no recrod in db
+	// We have not yet sent any ConfirmedInfoEvent and UnconfirmedInfoEvent, hence no recrod in db
+	assert.Equal(t, int64(0), overallStats.ActiveTvl)
 	assert.Equal(t, uint64(0), overallStats.UnconfirmedTvl)
+	assert.Equal(t, uint64(0), overallStats.PendingTvl)
 
 	// Test the top staker stats endpoint
 	stakerStats, _ := fetchStakerStatsEndpoint(t, testServer)
@@ -283,7 +284,6 @@ func TestStatsEndpoints(t *testing.T) {
 
 	expectedTvl := int64(activeEvents[0].StakingValue + activeEvents[1].StakingValue)
 	expectedTotalTvl := int64(expectedTvl) + int64(activeStakingEvent.StakingValue)
-	assert.Equal(t, expectedTvl, overallStats.ActiveTvl)
 	assert.Equal(t, expectedTotalTvl, overallStats.TotalTvl)
 	assert.Equal(t, int64(2), overallStats.ActiveDelegations)
 	assert.Equal(t, int64(3), overallStats.TotalDelegations)
@@ -310,6 +310,7 @@ func TestStatsEndpoints(t *testing.T) {
 
 	overallStats = fetchOverallStatsEndpoint(t, testServer)
 	assert.Equal(t, uint64(100), overallStats.UnconfirmedTvl)
+	assert.Equal(t, int64(90), overallStats.ActiveTvl)
 }
 
 func FuzzStatsEndpointReturnHighestUnconfirmedTvlFromEvents(f *testing.F) {
@@ -347,6 +348,8 @@ func FuzzStatsEndpointReturnHighestUnconfirmedTvlFromEvents(f *testing.F) {
 
 		overallStats = fetchOverallStatsEndpoint(t, testServer)
 		assert.Equal(t, &highestHeightEvent.UnconfirmedTvl, &overallStats.UnconfirmedTvl)
+		pendingTvl := highestHeightEvent.UnconfirmedTvl - highestHeightEvent.ConfirmedTvl
+		assert.Equal(t, pendingTvl, overallStats.PendingTvl)
 	})
 }