From 53b17c5431fac84ef012af8a9825011705b67d28 Mon Sep 17 00:00:00 2001
From: Brian Ginsburg <7957636+bgins@users.noreply.github.com>
Date: Mon, 18 Nov 2024 13:39:06 -0800
Subject: [PATCH] refactor: Use cleanup manager to unsubscribe web3
 subscriptions (#432)

* chore: Add unsubscribe helper

* refactor: Use cleanup manager to unsubscribe
---
 pkg/web3/events_jobcreator.go | 7 +------
 pkg/web3/events_mediation.go  | 7 +------
 pkg/web3/events_payments.go   | 6 +-----
 pkg/web3/events_pow.go        | 8 ++------
 pkg/web3/events_storage.go    | 7 +------
 pkg/web3/events_token.go      | 7 +------
 pkg/web3/utils.go             | 8 ++++++++
 7 files changed, 15 insertions(+), 35 deletions(-)

diff --git a/pkg/web3/events_jobcreator.go b/pkg/web3/events_jobcreator.go
index f9047223..d6bfa214 100644
--- a/pkg/web3/events_jobcreator.go
+++ b/pkg/web3/events_jobcreator.go
@@ -49,12 +49,7 @@ func (s *JobCreatorEventChannels) Start(
 	if err != nil {
 		return err
 	}
-
-	defer func() {
-		if jobAddedSub != nil {
-			jobAddedSub.Unsubscribe()
-		}
-	}()
+	cm.RegisterCallback(unsubscribeSub(jobAddedSub))
 
 	for {
 		select {
diff --git a/pkg/web3/events_mediation.go b/pkg/web3/events_mediation.go
index 5402f87d..43b511af 100644
--- a/pkg/web3/events_mediation.go
+++ b/pkg/web3/events_mediation.go
@@ -49,12 +49,7 @@ func (m *MediationEventChannels) Start(
 	if err != nil {
 		return err
 	}
-
-	defer func() {
-		if mediationRequestedSub != nil {
-			mediationRequestedSub.Unsubscribe()
-		}
-	}()
+	cm.RegisterCallback(unsubscribeSub(mediationRequestedSub))
 
 	for {
 		select {
diff --git a/pkg/web3/events_payments.go b/pkg/web3/events_payments.go
index f96303d9..61b5a6f2 100644
--- a/pkg/web3/events_payments.go
+++ b/pkg/web3/events_payments.go
@@ -49,12 +49,8 @@ func (p *PaymentEventChannels) Start(
 	if err != nil {
 		return err
 	}
+	cm.RegisterCallback(unsubscribeSub(paymentSub))
 
-	defer func() {
-		if paymentSub != nil {
-			paymentSub.Unsubscribe()
-		}
-	}()
 	for {
 		select {
 		case <-ctx.Done():
diff --git a/pkg/web3/events_pow.go b/pkg/web3/events_pow.go
index 79c9ccfc..49e2d6f4 100644
--- a/pkg/web3/events_pow.go
+++ b/pkg/web3/events_pow.go
@@ -3,6 +3,7 @@ package web3
 import (
 	"context"
 	"fmt"
+
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
 	"github.com/ethereum/go-ethereum/event"
 	"github.com/lilypad-tech/lilypad/pkg/system"
@@ -49,12 +50,7 @@ func (s *PowEventChannels) Start(
 	if err != nil {
 		return err
 	}
-
-	defer func() {
-		if newPowRoundSub != nil {
-			newPowRoundSub.Unsubscribe()
-		}
-	}()
+	cm.RegisterCallback(unsubscribeSub(newPowRoundSub))
 
 	for {
 		select {
diff --git a/pkg/web3/events_storage.go b/pkg/web3/events_storage.go
index 3d2076e1..3525bcaf 100644
--- a/pkg/web3/events_storage.go
+++ b/pkg/web3/events_storage.go
@@ -49,12 +49,7 @@ func (s *StorageEventChannels) Start(
 	if err != nil {
 		return err
 	}
-
-	defer func() {
-		if dealStateChangeSub != nil {
-			dealStateChangeSub.Unsubscribe()
-		}
-	}()
+	cm.RegisterCallback(unsubscribeSub(dealStateChangeSub))
 
 	for {
 		select {
diff --git a/pkg/web3/events_token.go b/pkg/web3/events_token.go
index 827c093e..1250f59e 100644
--- a/pkg/web3/events_token.go
+++ b/pkg/web3/events_token.go
@@ -52,12 +52,7 @@ func (t *TokenEventChannels) Start(
 	if err != nil {
 		return err
 	}
-
-	defer func() {
-		if transferSub != nil {
-			transferSub.Unsubscribe()
-		}
-	}()
+	cm.RegisterCallback(unsubscribeSub(transferSub))
 
 	for {
 		select {
diff --git a/pkg/web3/utils.go b/pkg/web3/utils.go
index 9c2d44e4..78bf934f 100644
--- a/pkg/web3/utils.go
+++ b/pkg/web3/utils.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/event"
 )
 
 func GetPublicKey(privateKey *ecdsa.PrivateKey) ecdsa.PublicKey {
@@ -68,3 +69,10 @@ func ConvertStringToInt64(st string) uint64 {
 	bigInt := ConvertStringToBigInt(st)
 	return bigInt.Uint64()
 }
+
+func unsubscribeSub(sub event.Subscription) func() error {
+	return func() error {
+		sub.Unsubscribe()
+		return nil
+	}
+}