From c1a17bc296e415ba33cf0b8243b2f0a7e607da4a Mon Sep 17 00:00:00 2001
From: Jesse de Wit <witdejesse@hotmail.com>
Date: Thu, 17 Nov 2022 17:11:51 +0100
Subject: [PATCH] A little prettier start and stop

---
 lnd_interceptor.go | 14 +++++++++++---
 main.go            | 35 ++++++++++++++++++++++++++++-------
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/lnd_interceptor.go b/lnd_interceptor.go
index 78389fe..45b9089 100644
--- a/lnd_interceptor.go
+++ b/lnd_interceptor.go
@@ -15,7 +15,8 @@ import (
 )
 
 type LndHtlcInterceptor struct {
-	client *LndClient
+	client        *LndClient
+	stopRequested bool
 }
 
 func NewLndHtlcInterceptor(client *LndClient) *LndHtlcInterceptor {
@@ -30,12 +31,16 @@ func (i *LndHtlcInterceptor) Start() error {
 	return i.intercept()
 }
 
-func (i *LndHtlcInterceptor) Stop() error {
-	return nil
+func (i *LndHtlcInterceptor) Stop() {
+	i.stopRequested = true
 }
 
 func (i *LndHtlcInterceptor) intercept() error {
 	for {
+		if i.stopRequested {
+			return nil
+		}
+
 		cancellableCtx, cancel := context.WithCancel(context.Background())
 		clientCtx := metadata.AppendToOutgoingContext(cancellableCtx, "macaroon", os.Getenv("LND_MACAROON_HEX"))
 		interceptorClient, err := client.routerClient.HtlcInterceptor(clientCtx)
@@ -47,6 +52,9 @@ func (i *LndHtlcInterceptor) intercept() error {
 		}
 
 		for {
+			if i.stopRequested {
+				return nil
+			}
 			request, err := interceptorClient.Recv()
 			if err != nil {
 				// If it is  just the error result of the context cancellation
diff --git a/main.go b/main.go
index 7c06105..1a89553 100644
--- a/main.go
+++ b/main.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"sync"
 
 	"github.com/btcsuite/btcd/btcec/v2"
 )
@@ -25,6 +26,7 @@ func main() {
 
 	client = NewLndClient()
 	interceptor := NewLndHtlcInterceptor(client)
+	s := NewGrpcServer()
 
 	info, err := client.GetInfo()
 	if err != nil {
@@ -37,16 +39,35 @@ func main() {
 		nodePubkey = info.Pubkey
 	}
 
-	go interceptor.Start()
-
 	go forwardingHistorySynchronize(client)
 	go channelsSynchronize(client)
 
-	s := NewGrpcServer()
-	err = s.Start()
-	if err != nil {
-		log.Fatalf("%v", err)
-	}
+	var wg sync.WaitGroup
+	wg.Add(2)
 
+	go func() {
+		err := interceptor.Start()
+		if err == nil {
+			log.Printf("Interceptor stopped.")
+		} else {
+			log.Printf("FATAL. Interceptor stopped with error: %v", err)
+		}
+		s.Stop()
+		wg.Done()
+	}()
+
+	go func() {
+		err := s.Start()
+		if err == nil {
+			log.Printf("GRPC server stopped.")
+		} else {
+			log.Printf("FATAL. GRPC server stopped with error: %v", err)
+		}
+
+		interceptor.Stop()
+		wg.Done()
+	}()
+
+	wg.Wait()
 	log.Printf("lspd exited")
 }