From a6f014c2c91e377e84bc8ddf3a0f02b95a7bb92b Mon Sep 17 00:00:00 2001
From: Shivam Sharma <shivam691999@gmail.com>
Date: Mon, 6 Dec 2021 14:25:33 +0530
Subject: [PATCH] fixes

---
 command/chain_watch.go    | 40 +++++++++++++++++++++++++++++----------
 command/server/server.go  |  2 --
 command/server/service.go |  8 ++++++--
 core/bor_events.go        |  6 ++++++
 4 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/command/chain_watch.go b/command/chain_watch.go
index 2ed37be9f6..05e52d2bc4 100644
--- a/command/chain_watch.go
+++ b/command/chain_watch.go
@@ -3,9 +3,13 @@ package main
 import (
 	"context"
 	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/ethereum/go-ethereum/command/flagset"
 	"github.com/ethereum/go-ethereum/command/server/proto"
+	"github.com/ethereum/go-ethereum/core"
 )
 
 // ChainWatchCommand is the command to group the peers commands
@@ -45,23 +49,39 @@ func (c *ChainWatchCommand) Run(args []string) int {
 
 	sub, err := borClt.ChainWatch(context.Background(), &proto.ChainWatchRequest{})
 	if err != nil {
-		panic(err)
+		c.UI.Error(err.Error())
+		return 1
 	}
 
+	signalCh := make(chan os.Signal, 1)
+	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)
+
+	go func() {
+		<-signalCh
+		sub.CloseSend()
+	}()
+
 	for {
 		msg, err := sub.Recv()
 		if err != nil {
 			// if err == EOF if finished on the other side
-			panic(err)
+			c.UI.Output(err.Error())
+			break
 		}
-		if msg.Type == "head" {
-			fmt.Println("Block Added : ", msg.Newchain)
-		} else if msg.Type == "fork" {
-			fmt.Println("New Fork Block :", msg.Newchain)
-		} else if msg.Type == "reorg" {
-			fmt.Println("Reorg Detected")
-			fmt.Println("Added :", msg.Newchain)
-			fmt.Println("Removed :", msg.Oldchain)
+		if msg.Type == core.Chain2HeadCanonicalEvent {
+			out := fmt.Sprintf("Block Added : %v", msg.Newchain)
+			c.UI.Output(out)
+		} else if msg.Type == core.Chain2HeadForkEvent {
+			out := fmt.Sprintf("New Fork Block : %v", msg.Newchain)
+			c.UI.Output(out)
+		} else if msg.Type == core.Chain2HeadReorgEvent {
+			c.UI.Output("Reorg Detected")
+
+			out := fmt.Sprintf("Added : %v", msg.Newchain)
+			c.UI.Output(out)
+
+			out = fmt.Sprintf("Removed : %v", msg.Oldchain)
+			c.UI.Output(out)
 		}
 
 		// fmt.Println(msg)
diff --git a/command/server/server.go b/command/server/server.go
index 8eb420d410..bab4673214 100644
--- a/command/server/server.go
+++ b/command/server/server.go
@@ -15,7 +15,6 @@ import (
 	"github.com/ethereum/go-ethereum/eth"
 	"github.com/ethereum/go-ethereum/eth/tracers"
 	"github.com/ethereum/go-ethereum/ethstats"
-	"github.com/ethereum/go-ethereum/event"
 	"github.com/ethereum/go-ethereum/graphql"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/metrics"
@@ -39,7 +38,6 @@ type Server struct {
 	backend    *eth.Ethereum
 	grpcServer *grpc.Server
 	tracer     *sdktrace.TracerProvider
-	headSub    event.Subscription
 }
 
 func NewServer(config *Config) (*Server, error) {
diff --git a/command/server/service.go b/command/server/service.go
index 4c01facf14..2a06460bbf 100644
--- a/command/server/service.go
+++ b/command/server/service.go
@@ -132,14 +132,18 @@ func (s *Server) ChainWatch(req *proto.ChainWatchRequest, reply proto.Bor_ChainW
 	chain2HeadChanSize := 10
 
 	chain2HeadCh := make(chan core.Chain2HeadEvent, chain2HeadChanSize)
-	s.headSub = s.backend.APIBackend.SubscribeChain2HeadEvent(chain2HeadCh)
+	headSub := s.backend.APIBackend.SubscribeChain2HeadEvent(chain2HeadCh)
+	defer headSub.Unsubscribe()
 
 	for {
 		msg := <-chain2HeadCh
 
-		reply.Send(&proto.ChainWatchResponse{Type: msg.Type,
+		err := reply.Send(&proto.ChainWatchResponse{Type: msg.Type,
 			Newchain: ConvertBlockToBlockStub(msg.NewChain),
 			Oldchain: ConvertBlockToBlockStub(msg.OldChain),
 		})
+		if err != nil {
+			return err
+		}
 	}
 }
diff --git a/core/bor_events.go b/core/bor_events.go
index e47afba1a5..68f916fe79 100644
--- a/core/bor_events.go
+++ b/core/bor_events.go
@@ -9,6 +9,12 @@ type StateSyncEvent struct {
 	Data *types.StateSyncData
 }
 
+var (
+	Chain2HeadReorgEvent     = "reorg"
+	Chain2HeadCanonicalEvent = "head"
+	Chain2HeadForkEvent      = "fork"
+)
+
 // For tracking reorgs related information
 type Chain2HeadEvent struct {
 	NewChain []*types.Block