From 5f907ca1e6ffa041e919472c1ebfe216ad9d0dac Mon Sep 17 00:00:00 2001
From: Ian Norden <iannordenn@gmail.com>
Date: Wed, 19 Aug 2020 00:57:57 -0500
Subject: [PATCH] goimports + minor fixes

---
 pkg/api.go     |  7 +------
 pkg/reader.go  | 12 ++++++------
 pkg/service.go | 23 +++++++++++++----------
 3 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/pkg/api.go b/pkg/api.go
index 666e6ab..57d8c1d 100644
--- a/pkg/api.go
+++ b/pkg/api.go
@@ -17,15 +17,10 @@ package statediff
 
 import (
 	"context"
+
 	"github.com/ethereum/go-ethereum/statediff"
 )
 
-// APIName is the namespace used for the state diffing service API
-const APIName = "statediff"
-
-// APIVersion is the version of the state diffing service API
-const APIVersion = "0.0.1"
-
 // PublicStateDiffAPI provides an RPC interface
 // that can be used to fetch historical diffs from leveldb directly
 type PublicStateDiffAPI struct {
diff --git a/pkg/reader.go b/pkg/reader.go
index b4792d0..95a7bc6 100644
--- a/pkg/reader.go
+++ b/pkg/reader.go
@@ -29,9 +29,9 @@ import (
 
 // LvlDBReader exposes the necessary read functions on lvldb
 type LvlDBReader struct {
-	ethDB         ethdb.Database
-	stateDB       state.Database
-	chainConfig   *params.ChainConfig
+	ethDB       ethdb.Database
+	stateDB     state.Database
+	chainConfig *params.ChainConfig
 }
 
 // NewLvlDBReader creates a new LvlDBReader
@@ -41,8 +41,8 @@ func NewLvlDBReader(path, ancient string, chainConfig *params.ChainConfig) (*Lvl
 		return nil, err
 	}
 	return &LvlDBReader{
-		ethDB: edb,
-		stateDB: state.NewDatabase(edb),
+		ethDB:       edb,
+		stateDB:     state.NewDatabase(edb),
 		chainConfig: chainConfig,
 	}, nil
 }
@@ -98,4 +98,4 @@ func (ldr *LvlDBReader) GetTdByHash(hash common.Hash) (*big.Int, error) {
 // StateDB returns the underlying statedb
 func (ldr *LvlDBReader) StateDB() state.Database {
 	return ldr.stateDB
-}
\ No newline at end of file
+}
diff --git a/pkg/service.go b/pkg/service.go
index e9ad044..db46b7b 100644
--- a/pkg/service.go
+++ b/pkg/service.go
@@ -19,6 +19,7 @@ import (
 	"bytes"
 	"fmt"
 	"math/big"
+	"sync"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/state"
@@ -45,7 +46,7 @@ type IService interface {
 	// APIs(), Protocols(), Start() and Stop()
 	node.Service
 	// Main event loop for processing state diffs
-	Loop()
+	Loop(wg *sync.WaitGroup)
 	// Method to get state diff object at specific block
 	StateDiffAt(blockNumber uint64, params statediff.Params) (*statediff.Payload, error)
 	// Method to get state trie object at specific block
@@ -65,9 +66,9 @@ type Service struct {
 // NewStateDiffService creates a new statediff.Service
 func NewStateDiffService(lvlDBReader lvlDBReader) (*Service, error) {
 	return &Service{
-		lvlDBReader:       lvlDBReader,
-		Builder:           statediff.NewBuilder(lvlDBReader.StateDB()),
-		QuitChan:          make(chan bool),
+		lvlDBReader: lvlDBReader,
+		Builder:     statediff.NewBuilder(lvlDBReader.StateDB()),
+		QuitChan:    make(chan bool),
 	}, nil
 }
 
@@ -89,11 +90,13 @@ func (sds *Service) APIs() []rpc.API {
 }
 
 // Loop is an empty service loop for awaiting rpc requests
-func (sds *Service) Loop() {
+func (sds *Service) Loop(wg *sync.WaitGroup) {
+	wg.Add(1)
 	for {
 		select {
 		case <-sds.QuitChan:
-			log.Info("Closing the statediff service")
+			log.Info("closing the statediff service loop")
+			wg.Done()
 			return
 		}
 	}
@@ -148,7 +151,7 @@ func (sds *Service) newPayload(stateObject []byte, block *types.Block, params st
 	}
 	if params.IncludeTD {
 		var err error
-		payload.TotalDifficulty , err = sds.lvlDBReader.GetTdByHash(block.Hash())
+		payload.TotalDifficulty, err = sds.lvlDBReader.GetTdByHash(block.Hash())
 		if err != nil {
 			return nil, err
 		}
@@ -192,14 +195,14 @@ func (sds *Service) processStateTrie(block *types.Block, params statediff.Params
 
 // Start is used to begin the service
 func (sds *Service) Start(*p2p.Server) error {
-	log.Info("Starting statediff service")
-	go sds.Loop()
+	log.Info("starting statediff service")
+	go sds.Loop(new(sync.WaitGroup))
 	return nil
 }
 
 // Stop is used to close down the service
 func (sds *Service) Stop() error {
-	log.Info("Stopping statediff service")
+	log.Info("stopping statediff service")
 	close(sds.QuitChan)
 	return nil
 }