From 5402853976420b6120a52cf2d60306e4b9217356 Mon Sep 17 00:00:00 2001
From: Jeffrey Picard <jeff@jeffreypicard.com>
Date: Mon, 7 Mar 2022 20:25:56 +0000
Subject: [PATCH] signal handling and cleanup functions

---
 db/db.go         | 19 ++++++++++++++-----
 db/db_test.go    | 21 ++++-----------------
 main.go          | 27 ++++++++++++---------------
 server/server.go |  5 ++++-
 4 files changed, 34 insertions(+), 38 deletions(-)

diff --git a/db/db.go b/db/db.go
index 4c3c863..ad4dfe2 100644
--- a/db/db.go
+++ b/db/db.go
@@ -448,7 +448,7 @@ func GetWriteDBCF(name string) (*grocksdb.DB, []*grocksdb.ColumnFamilyHandle, er
 }
 
 // GetProdDB returns a db that is used for production.
-func GetProdDB(name string, secondaryPath string) (*ReadOnlyDBColumnFamily, error) {
+func GetProdDB(name string, secondaryPath string) (*ReadOnlyDBColumnFamily, func(), error) {
 	prefixNames := prefixes.GetPrefixes()
 	// additional prefixes that aren't in the code explicitly
 	cfNames := []string{"default", "e", "d", "c"}
@@ -458,11 +458,20 @@ func GetProdDB(name string, secondaryPath string) (*ReadOnlyDBColumnFamily, erro
 	}
 
 	db, err := GetDBColumnFamlies(name, secondaryPath, cfNames)
-	if err != nil {
-		return nil, err
+
+	cleanup := func() {
+		db.DB.Close()
+		err = os.RemoveAll(fmt.Sprintf("./%s", secondaryPath))
+		if err != nil {
+			log.Println(err)
+		}
 	}
 
-	return db, nil
+	if err != nil {
+		return nil, cleanup, err
+	}
+
+	return db, cleanup, nil
 }
 
 func GetDBColumnFamlies(name string, secondayPath string, cfNames []string) (*ReadOnlyDBColumnFamily, error) {
@@ -632,7 +641,7 @@ func DetectChanges(db *ReadOnlyDBColumnFamily) error {
 
 	if db.LastState == nil || lastHeight < state.Height {
 		for height := lastHeight + 1; height <= state.Height; height++ {
-			log.Info("advancing to", height)
+			log.Info("advancing to: ", height)
 			Advance(db, height)
 		}
 		//TODO: ClearCache
diff --git a/db/db_test.go b/db/db_test.go
index 4e2491b..f380154 100644
--- a/db/db_test.go
+++ b/db/db_test.go
@@ -249,14 +249,8 @@ func TestCatFullDB(t *testing.T) {
 	dbPath := "/mnt/sda/wallet_server/_data/lbry-rocksdb/"
 	// dbPath := "/mnt/d/data/snapshot_1072108/lbry-rocksdb/"
 	secondaryPath := "asdf"
-	db, err := dbpkg.GetProdDB(dbPath, secondaryPath)
-	toDefer := func() {
-		db.DB.Close()
-		err = os.RemoveAll("./asdf")
-		if err != nil {
-			log.Println(err)
-		}
-	}
+	db, toDefer, err := dbpkg.GetProdDB(dbPath, secondaryPath)
+
 	defer toDefer()
 	if err != nil {
 		t.Error(err)
@@ -283,18 +277,11 @@ func TestOpenFullDB(t *testing.T) {
 	// url := "lbry://@lbry"
 	// url := "lbry://@lbry#3fda836a92faaceedfe398225fb9b2ee2ed1f01a"
 	// url := "lbry://@lbry$1"
-	url := "lbry://@lothrop:2/lothrop-livestream-games-and-code:c"
+	url := "https://lbry.tv/@lothrop:2/lothrop-livestream-games-and-code:c"
 	dbPath := "/mnt/sda/wallet_server/_data/lbry-rocksdb/"
 	// dbPath := "/mnt/d/data/snapshot_1072108/lbry-rocksdb/"
 	secondaryPath := "asdf"
-	db, err := dbpkg.GetProdDB(dbPath, secondaryPath)
-	toDefer := func() {
-		db.DB.Close()
-		err = os.RemoveAll("./asdf")
-		if err != nil {
-			log.Println(err)
-		}
-	}
+	db, toDefer, err := dbpkg.GetProdDB(dbPath, secondaryPath)
 	defer toDefer()
 	if err != nil {
 		t.Error(err)
diff --git a/main.go b/main.go
index 262ff01..6936093 100644
--- a/main.go
+++ b/main.go
@@ -33,27 +33,24 @@ func main() {
 		defer cancel()
 
 		// TODO: Figure out if / where we need signal handling
-		// interrupt := make(chan os.Signal, 1)
-		// signal.Notify(interrupt, os.Interrupt, syscall.SIGTERM)
-		// defer signal.Stop(interrupt)
+
+		initsignals()
+		interrupt := interruptListener()
 
 		s := server.MakeHubServer(ctxWCancel, args)
-		s.Run()
+		go s.Run()
 
-		// select {
-		// case <-interrupt:
-		// 	break
-		// case <-ctx.Done():
-		// 	break
-		// }
+		defer func() {
+			log.Println("Shutting down server...")
 
-		// log.Println("Shutting down server...")
+			s.EsClient.Stop()
+			s.GrpcServer.GracefulStop()
+			s.DBCleanup()
 
-		// s.EsClient.Stop()
-		// s.GrpcServer.GracefulStop()
-
-		// log.Println("Returning from main...")
+			log.Println("Returning from main...")
+		}()
 
+		<-interrupt
 		return
 	} else if args.CmdType == server.DBCmd {
 		options := &db.IterOptions{
diff --git a/server/server.go b/server/server.go
index f6deb7b..d02051c 100644
--- a/server/server.go
+++ b/server/server.go
@@ -44,6 +44,7 @@ type Server struct {
 	PeerSubsMut      sync.RWMutex
 	NumPeerSubs      *int64
 	ExternalIP       net.IP
+	DBCleanup        func()
 	pb.UnimplementedHubServer
 }
 
@@ -188,8 +189,9 @@ func MakeHubServer(ctx context.Context, args *Args) *Server {
 
 	//TODO: is this the right place to load the db?
 	var myDB *db.ReadOnlyDBColumnFamily
+	var dbCleanup = func() {}
 	if !args.DisableResolve {
-		myDB, err = db.GetProdDB(args.DBPath, "readonlytmp")
+		myDB, dbCleanup, err = db.GetProdDB(args.DBPath, "readonlytmp")
 		if err != nil {
 			// Can't load the db, fail loudly
 			log.Fatalln(err)
@@ -215,6 +217,7 @@ func MakeHubServer(ctx context.Context, args *Args) *Server {
 		PeerSubsMut:      sync.RWMutex{},
 		NumPeerSubs:      numSubs,
 		ExternalIP:       net.IPv4(127, 0, 0, 1),
+		DBCleanup:        dbCleanup,
 	}
 
 	// Start up our background services