Made the healthcheck more robust and added some recovery mechanism

api/action.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package api
 
 import (

api/asset.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package api
 
 import (

api/block.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package api
 
 import (

api/genesis.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package api
 
 import (

api/health.go

@@ -1,20 +1,61 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package api
 
 import (
+	"database/sql"
+	"log"
+	"net"
 	"net/http"
+	"time"
 
 	"github.com/gin-gonic/gin"
 )
 
-// HealthCheck performs a comprehensive health check of the subscriber
-func HealthCheck() gin.HandlerFunc {
+// checkGRPC checks the reachability of the gRPC server.
+func checkGRPC(grpcPort string) string {
+	conn, err := net.DialTimeout("tcp", grpcPort, 2*time.Second)
+	if err != nil {
+		log.Printf("gRPC connection failed: %v", err)
+		return "unreachable"
+	}
+	defer conn.Close()
+	return "reachable"
+}
+
+// checkDatabase checks the reachability of the database.
+func checkDatabase(db *sql.DB) string {
+	if err := db.Ping(); err != nil {
+		log.Printf("Database connection failed: %v", err)
+		return "unreachable"
+	}
+	return "reachable"
+}
+
+// HealthCheck performs a comprehensive health check of the subscriber.
+func HealthCheck(grpcPort string, db *sql.DB) gin.HandlerFunc {
 	return func(c *gin.Context) {
-		// If all checks pass
-		c.JSON(http.StatusOK, gin.H{
-			"status": "ok",
+		grpcStatus := checkGRPC(grpcPort)
+		databaseStatus := checkDatabase(db)
+
+		// Determine the overall status
+		status := "ok"
+		if grpcStatus == "unreachable" || databaseStatus == "unreachable" {
+			status = "service unavailable"
+		}
+
+		// Return the health status
+		httpStatusCode := http.StatusOK
+		if status == "service unavailable" {
+			httpStatusCode = http.StatusServiceUnavailable
+		}
+
+		c.JSON(httpStatusCode, gin.H{
+			"status": status,
 			"details": gin.H{
-				"database": "reachable",
-				"grpc":     "reachable",
+				"database": databaseStatus,
+				"grpc":     grpcStatus,
 			},
 		})
 	}

api/transaction.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package api
 
 import (

config/config.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package config
 
 import (

consts/action_names.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package consts
 
 import "github.com/nuklai/nuklaivm/consts"

db/db.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package db
 
 import (

grpc/server.go

@@ -1,9 +1,13 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package grpc
 
 import (
 	"context"
 	"database/sql"
 	"encoding/json"
+	"errors"
 	"log"
 	"net"
 	"strings"
@@ -31,8 +35,14 @@ type Server struct {
 	parser chain.Parser
 }
 
-// startGRPCServer starts the gRPC server for receiving block data
-func StartGRPCServer(db *sql.DB, port string) {
+// StartGRPCServer starts the gRPC server for receiving block data
+func StartGRPCServer(db *sql.DB, port string) error {
+	defer func() {
+		if r := recover(); r != nil {
+			log.Printf("Recovered from panic in StartGRPCServer: %v", r)
+		}
+	}()
+
 	// Load the whitelist
 	LoadWhitelist()
 
@@ -43,7 +53,8 @@ func StartGRPCServer(db *sql.DB, port string) {
 
 	lis, err := net.Listen("tcp", port)
 	if err != nil {
-		log.Fatalf("Failed to listen on port %s: %v", port, err)
+		log.Printf("Failed to listen on port %s: %v", port, err)
+		return err
 	}
 
 	// Use insecure credentials to allow plaintext communication
@@ -60,14 +71,32 @@ func StartGRPCServer(db *sql.DB, port string) {
 	reflection.Register(grpcServer)
 
 	log.Printf("External Subscriber server is listening on port %s...\n", port)
-	if err := grpcServer.Serve(lis); err != nil {
-		log.Fatalf("Failed to serve: %v", err)
+	return grpcServer.Serve(lis)
+}
+
+// StartGRPCServerWithRetries retries gRPC server startup in case of failure
+func StartGRPCServerWithRetries(db *sql.DB, port string, retries int) {
+	for i := 0; i < retries; i++ {
+		err := StartGRPCServer(db, port)
+		if err != nil {
+			log.Printf("gRPC server failed to start: %v. Retrying (%d/%d)...", err, i+1, retries)
+			time.Sleep(5 * time.Second)
+			continue
+		}
+		return // Successful start
 	}
+	log.Fatal("gRPC server failed to start after maximum retries")
 }
 
 // Initialize receives genesis data for initialization and saves it to the database
 // Initialize receives genesis data for initialization and saves it to the database
 func (s *Server) Initialize(ctx context.Context, req *pb.InitializeRequest) (*emptypb.Empty, error) {
+	defer func() {
+		if r := recover(); r != nil {
+			log.Printf("Recovered from panic in Initialize: %v", r)
+		}
+	}()
+
 	log.Println("Initializing External Subscriber with genesis data...")
 
 	// Decode genesis data
@@ -108,6 +137,17 @@ func (s *Server) Initialize(ctx context.Context, req *pb.InitializeRequest) (*em
 
 // AcceptBlock processes a new block, saves relevant data to the database, and stores transactions and actions
 func (s *Server) AcceptBlock(ctx context.Context, req *pb.BlockRequest) (*emptypb.Empty, error) {
+	defer func() {
+		if r := recover(); r != nil {
+			log.Printf("Recovered from panic in AcceptBlock: %v", r)
+		}
+	}()
+
+	if s.parser == nil {
+		log.Println("Parser is not initialized. Rejecting the request.")
+		return nil, errors.New("parser not initialized")
+	}
+
 	blockData := req.GetBlockData()
 
 	// Attempt to unmarshal the executed block using UnmarshalExecutedBlock

grpc/whitelist.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package grpc
 
 import (

main.go

@@ -1,4 +1,6 @@
-// Additional API endpoints to retrieve blockchain-related information from the database
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package main
 
 import (
@@ -25,10 +27,16 @@ func main() {
 	defer database.Close()
 
 	// Start the gRPC server
-	go grpc.StartGRPCServer(database, "50051")
+	grpcPort := "50051"
+	go grpc.StartGRPCServerWithRetries(database, grpcPort, 60)
 
 	// Setup Gin router
-	r := gin.Default()
+	gin.SetMode(gin.ReleaseMode)
+
+	r := gin.New()
+	r.Use(gin.Logger(), gin.Recovery())
+
+	r.SetTrustedProxies(nil)
 
 	// Add CORS middleware
 	r.Use(cors.New(cors.Config{
@@ -40,7 +48,7 @@ func main() {
 	}))
 
 	// Health endpoint
-	r.GET("/health", api.HealthCheck())
+	r.GET("/health", api.HealthCheck(":"+grpcPort, database))
 
 	// Other endpoints
 	r.GET("/genesis", api.GetGenesisData(database))

models/action.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package models
 
 import (

models/asset.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package models
 
 import (

models/block.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package models
 
 import (

models/transaction.go

@@ -1,3 +1,6 @@
+// Copyright (C) 2024, Nuklai. All rights reserved.
+// See the file LICENSE for licensing terms.
+
 package models
 
 import (