From 3b572176aa1474aba01203bd2b37b20ae7821020 Mon Sep 17 00:00:00 2001
From: bruwbird <bruwbird@mercari.com>
Date: Sun, 22 Sep 2024 11:17:49 +0900
Subject: [PATCH] glightning: add shutdown subscription

Adds a new shutdown subscription to the plugin system,
allowing plugins to be notified when lightningd is shutting down.
This enables plugins to perform any necessary cleanup
or finalization tasks before the daemon exits.
---
 glightning/plugin.go      | 45 ++++++++++++++++++++++++++++++++-------
 glightning/plugin_test.go | 26 ++++++++++++++++++++++
 2 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/glightning/plugin.go b/glightning/plugin.go
index 47c9dd9..da111b9 100644
--- a/glightning/plugin.go
+++ b/glightning/plugin.go
@@ -26,6 +26,7 @@ const (
 	_Forward        Subscription = "forward_event"
 	_SendPaySuccess Subscription = "sendpay_success"
 	_SendPayFailure Subscription = "sendpay_failure"
+	_Shutdown       Subscription = "shutdown"
 	_PeerConnected  Hook         = "peer_connected"
 	_DbWrite        Hook         = "db_write"
 	_InvoicePayment Hook         = "invoice_payment"
@@ -38,7 +39,7 @@ const (
 var lightningMethodRegistry map[string]*jrpc2.Method
 
 // The custommsg plugin hook is the receiving counterpart to the dev-sendcustommsg RPC method
-//and allows plugins to handle messages that are not handled internally.
+// and allows plugins to handle messages that are not handled internally.
 type CustomMsgReceivedEvent struct {
 	PeerId  string `json:"peer_id"`
 	Payload string `json:"payload"`
@@ -81,8 +82,9 @@ func (pc *CustomMsgReceivedEvent) Fail() *CustomMsgReceivedResponse {
 }
 
 // This hook is called whenever a peer has connected and successfully completed
-//   the cryptographic handshake. The parameters have the following structure if
-//   there is a channel with the peer:
+//
+//	the cryptographic handshake. The parameters have the following structure if
+//	there is a channel with the peer:
 type PeerConnectedEvent struct {
 	Peer PeerEvent `json:"peer"`
 	hook func(*PeerConnectedEvent) (*PeerConnectedResponse, error)
@@ -450,10 +452,11 @@ func (rc *RpcCommandEvent) ReturnError(errMsg string, errCode int) (*RpcCommandR
 // its result determines how `lightningd` should treat that HTLC.
 //
 // Warning: `lightningd` will replay the HTLCs for which it doesn't have a final
-//   verdict during startup. This means that, if the plugin response wasn't
-//   processed before the HTLC was forwarded, failed, or resolved, then the plugin
-//   may see the same HTLC again during startup. It is therefore paramount that the
-//   plugin is idempotent if it talks to an external system.
+//
+//	verdict during startup. This means that, if the plugin response wasn't
+//	processed before the HTLC was forwarded, failed, or resolved, then the plugin
+//	may see the same HTLC again during startup. It is therefore paramount that the
+//	plugin is idempotent if it talks to an external system.
 type HtlcAcceptedEvent struct {
 	Onion Onion     `json:"onion"`
 	Htlc  HtlcOffer `json:"htlc"`
@@ -763,6 +766,25 @@ func (e *WarnEvent) Call() (jrpc2.Result, error) {
 	return nil, nil
 }
 
+type ShutdownEvent struct {
+	cb func()
+}
+
+func (e *ShutdownEvent) Name() string {
+	return string(_Shutdown)
+}
+
+func (e *ShutdownEvent) New() interface{} {
+	return &ShutdownEvent{
+		cb: e.cb,
+	}
+}
+
+func (e *ShutdownEvent) Call() (jrpc2.Result, error) {
+	e.cb()
+	return nil, nil
+}
+
 type OptionType string
 
 const _String OptionType = "string"
@@ -1182,7 +1204,8 @@ func (p *Plugin) Log(message string, level LogLevel) {
 }
 
 // Map for registering hooks. Not the *most* elegant but
-//   it'll do for now.
+//
+//	it'll do for now.
 type Hooks struct {
 	PeerConnected     func(*PeerConnectedEvent) (*PeerConnectedResponse, error)
 	DbWrite           func(*DbWriteEvent) (*DbWriteResponse, error)
@@ -1559,6 +1582,12 @@ func (p *Plugin) SubscribeForwardings(cb func(c *Forwarding)) {
 	})
 }
 
+func (p *Plugin) SubscribeShutdown(cb func()) {
+	p.subscribe(&ShutdownEvent{
+		cb: cb,
+	})
+}
+
 func (p *Plugin) subscribe(subscription jrpc2.ServerMethod) {
 	p.server.Register(subscription)
 	p.subscriptions = append(p.subscriptions, subscription.Name())
diff --git a/glightning/plugin_test.go b/glightning/plugin_test.go
index d9d7b88..5e62ad5 100644
--- a/glightning/plugin_test.go
+++ b/glightning/plugin_test.go
@@ -756,6 +756,32 @@ func TestSubscription_Disconnected(t *testing.T) {
 	runTest(t, plugin, msg+"\n\n", "")
 }
 
+func TestSubscription_Shutdown(t *testing.T) {
+	var wg sync.WaitGroup
+	defer await(t, &wg)
+	shutdownCalled := make(chan struct{})
+
+	wg.Add(1)
+	initFn := getInitFunc(t, func(t *testing.T, options map[string]glightning.Option, config *glightning.Config) {
+		t.Error("Should not have called init when calling get manifest")
+	})
+	plugin := glightning.NewPlugin(initFn)
+	plugin.SubscribeShutdown(func() {
+		defer wg.Done()
+		shutdownCalled <- struct{}{}
+	})
+
+	msg := `{"jsonrpc":"2.0","method":"shutdown"}`
+
+	runTest(t, plugin, msg+"\n\n", "")
+
+	select {
+	case <-shutdownCalled:
+	case <-time.After(1 * time.Second):
+		t.Fatal("SubscribeShutdown was not called")
+	}
+}
+
 func await(t *testing.T, wg *sync.WaitGroup) {
 	awaitWithTimeout(t, wg, 1*time.Second)
 }