From 8de8a6815c9f3dd349c89f98a97376d4342d4d85 Mon Sep 17 00:00:00 2001
From: Krishna Upadhyaya <krishnau1604@gmail.com>
Date: Tue, 26 Jul 2022 18:08:25 +0530
Subject: [PATCH] Integrated heimdall gRPC server flow

---
 cmd/utils/bor_flags.go                   |  2 +-
 consensus/bor/heimdallgrpc/client.go     | 18 ++++++++----------
 consensus/bor/heimdallgrpc/state_sync.go |  2 +-
 eth/ethconfig/config.go                  | 10 +++++++++-
 internal/cli/server/config.go            |  2 +-
 5 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/cmd/utils/bor_flags.go b/cmd/utils/bor_flags.go
index 6369b251c6..256df31e97 100644
--- a/cmd/utils/bor_flags.go
+++ b/cmd/utils/bor_flags.go
@@ -35,7 +35,7 @@ var (
 	HeimdallgRPCAddressFlag = cli.StringFlag{
 		Name:  "bor.heimdallgRPC",
 		Usage: "Address of Heimdall gRPC service",
-		Value: ":3132",
+		Value: "",
 	}
 
 	// BorFlags all bor related flags
diff --git a/consensus/bor/heimdallgrpc/client.go b/consensus/bor/heimdallgrpc/client.go
index a52fc20dca..caf7b14002 100644
--- a/consensus/bor/heimdallgrpc/client.go
+++ b/consensus/bor/heimdallgrpc/client.go
@@ -1,8 +1,7 @@
 package heimdallgrpc
 
 import (
-	"log"
-
+	"github.com/ethereum/go-ethereum/log"
 	proto "github.com/maticnetwork/polyproto/heimdall"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
@@ -13,25 +12,24 @@ const (
 )
 
 type HeimdallGRPCClient struct {
-	conn    *grpc.ClientConn
-	client  proto.HeimdallClient
-	closeCh chan struct{}
+	conn   *grpc.ClientConn
+	client proto.HeimdallClient
 }
 
 func NewHeimdallGRPCClient(address string) *HeimdallGRPCClient {
 	conn, err := grpc.Dial(address, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
-		log.Fatalf("Failed to connect to Heimdall gRPC server: %v", err)
+		log.Error("Failed to connect to Heimdall gRPC server: %v", err)
+		panic(err)
 	}
 
 	return &HeimdallGRPCClient{
-		conn:    conn,
-		client:  proto.NewHeimdallClient(conn),
-		closeCh: make(chan struct{}),
+		conn:   conn,
+		client: proto.NewHeimdallClient(conn),
 	}
 }
 
 func (h *HeimdallGRPCClient) Close() {
-	close(h.closeCh)
+	log.Debug("Shutdown detected, Closing Heimdall gRPC client")
 	h.conn.Close()
 }
diff --git a/consensus/bor/heimdallgrpc/state_sync.go b/consensus/bor/heimdallgrpc/state_sync.go
index aaa4736548..2f36895687 100644
--- a/consensus/bor/heimdallgrpc/state_sync.go
+++ b/consensus/bor/heimdallgrpc/state_sync.go
@@ -9,7 +9,7 @@ import (
 	proto "github.com/maticnetwork/polyproto/heimdall"
 )
 
-func (h *HeimdallGRPCClient) StateSyncEvents(fromID uint64, to int64) ([]*clerk.EventRecordWithTime, error) {
+func (h *HeimdallGRPCClient) StateSyncEvents(ctx context.Context, fromID uint64, to int64) ([]*clerk.EventRecordWithTime, error) {
 	eventRecords := make([]*clerk.EventRecordWithTime, 0)
 
 	req := &proto.StateSyncEventsRequest{
diff --git a/eth/ethconfig/config.go b/eth/ethconfig/config.go
index 10616b60c5..adb36ffadd 100644
--- a/eth/ethconfig/config.go
+++ b/eth/ethconfig/config.go
@@ -32,6 +32,7 @@ import (
 	"github.com/ethereum/go-ethereum/consensus/bor/contract"
 	"github.com/ethereum/go-ethereum/consensus/bor/heimdall" //nolint:typecheck
 	"github.com/ethereum/go-ethereum/consensus/bor/heimdall/span"
+	"github.com/ethereum/go-ethereum/consensus/bor/heimdallgrpc"
 	"github.com/ethereum/go-ethereum/consensus/clique"
 	"github.com/ethereum/go-ethereum/consensus/ethash"
 	"github.com/ethereum/go-ethereum/core"
@@ -249,7 +250,14 @@ func CreateConsensusEngine(stack *node.Node, chainConfig *params.ChainConfig, et
 		if ethConfig.WithoutHeimdall {
 			return bor.New(chainConfig, db, blockchainAPI, spanner, nil, genesisContractsClient)
 		} else {
-			return bor.New(chainConfig, db, blockchainAPI, spanner, heimdall.NewHeimdallClient(ethConfig.HeimdallURL), genesisContractsClient)
+			var heimdallClient bor.IHeimdallClient
+			if ethConfig.HeimdallgRPCAddress != "" {
+				heimdallClient = heimdallgrpc.NewHeimdallGRPCClient(ethConfig.HeimdallgRPCAddress)
+			} else {
+				heimdallClient = heimdall.NewHeimdallClient(ethConfig.HeimdallURL)
+			}
+
+			return bor.New(chainConfig, db, blockchainAPI, spanner, heimdallClient, genesisContractsClient)
 		}
 	} else {
 		switch config.PowMode {
diff --git a/internal/cli/server/config.go b/internal/cli/server/config.go
index 522d4004ae..37caa4fcc2 100644
--- a/internal/cli/server/config.go
+++ b/internal/cli/server/config.go
@@ -414,7 +414,7 @@ func DefaultConfig() *Config {
 		Heimdall: &HeimdallConfig{
 			URL:         "http://localhost:1317",
 			Without:     false,
-			GRPCAddress: ":3132",
+			GRPCAddress: "",
 		},
 		SyncMode: "full",
 		GcMode:   "full",