From c474956da85e0d5a18d12b3fe63de53bb544a856 Mon Sep 17 00:00:00 2001
From: maskpp <maskpp266@gmail.com>
Date: Sun, 14 Jan 2024 10:21:05 +0800
Subject: [PATCH] fix(rpc): fix a bug / update logic (#501)

Co-authored-by: David <david@taiko.xyz>
---
 common/utils/utils.go |  8 ++++++++
 driver/driver.go      |  1 +
 go.mod                |  1 +
 go.sum                |  2 ++
 pkg/rpc/methods.go    | 24 ++++++++++--------------
 pkg/rpc/utils.go      | 13 +++++++------
 6 files changed, 29 insertions(+), 20 deletions(-)
 create mode 100644 common/utils/utils.go

diff --git a/common/utils/utils.go b/common/utils/utils.go
new file mode 100644
index 000000000..5dceaffdc
--- /dev/null
+++ b/common/utils/utils.go
@@ -0,0 +1,8 @@
+package utils
+
+import "github.com/modern-go/reflect2"
+
+// IsNil checks if the interface is empty.
+func IsNil(i interface{}) bool {
+	return i == nil || reflect2.IsNil(i)
+}
diff --git a/driver/driver.go b/driver/driver.go
index a32c6c30a..119535ec3 100644
--- a/driver/driver.go
+++ b/driver/driver.go
@@ -122,6 +122,7 @@ func (d *Driver) Start() error {
 
 // Close closes the driver instance.
 func (d *Driver) Close(ctx context.Context) {
+	d.l1HeadSub.Unsubscribe()
 	d.state.Close()
 	d.wg.Wait()
 }
diff --git a/go.mod b/go.mod
index 63be10569..e81e15dac 100644
--- a/go.mod
+++ b/go.mod
@@ -9,6 +9,7 @@ require (
 	github.com/ethereum/go-ethereum v1.13.8
 	github.com/go-resty/resty/v2 v2.7.0
 	github.com/labstack/echo/v4 v4.11.1
+	github.com/modern-go/reflect2 v1.0.2
 	github.com/phayes/freeport v0.0.0-20220201140144-74d24b5ae9f5
 	github.com/prysmaticlabs/prysm/v4 v4.0.1
 	github.com/stretchr/testify v1.8.4
diff --git a/go.sum b/go.sum
index 82aa72a3a..7f796a180 100644
--- a/go.sum
+++ b/go.sum
@@ -299,6 +299,8 @@ github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJ
 github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
 github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
+github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
+github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
 github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 h1:RWengNIwukTxcDr9M+97sNutRR1RKhG96O6jWumTTnw=
 github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826/go.mod h1:TaXosZuwdSHYgviHp1DAtfrULt5eUgsSMsZf+YrPgl8=
 github.com/moul/http2curl v1.0.0/go.mod h1:8UbvGypXm98wA/IqH45anm5Y2Z6ep6O31QGOAZ3H0fQ=
diff --git a/pkg/rpc/methods.go b/pkg/rpc/methods.go
index a875c0dc8..27ff54d65 100644
--- a/pkg/rpc/methods.go
+++ b/pkg/rpc/methods.go
@@ -345,21 +345,17 @@ func (c *Client) GetProtocolStateVariables(opts *bind.CallOpts) (*struct {
 	A bindings.TaikoDataSlotA
 	B bindings.TaikoDataSlotB
 }, error) {
-	var (
-		ctxWithTimeout context.Context
-		cancel         context.CancelFunc
-	)
-	if opts != nil && opts.Context != nil {
-		if _, ok := opts.Context.Deadline(); !ok {
-			ctxWithTimeout, cancel = context.WithTimeout(opts.Context, defaultWaitReceiptTimeout)
-			defer cancel()
-			opts.Context = ctxWithTimeout
-		}
-	} else {
-		ctxWithTimeout, cancel = context.WithTimeout(context.Background(), defaultWaitReceiptTimeout)
-		defer cancel()
-		opts = &bind.CallOpts{Context: ctxWithTimeout}
+	if opts == nil {
+		opts = &bind.CallOpts{}
+	}
+
+	var ctx = context.Background()
+	if opts.Context != nil {
+		ctx = opts.Context
 	}
+	ctxWithTimeout, cancel := context.WithTimeout(ctx, defaultWaitReceiptTimeout)
+	defer cancel()
+	opts.Context = ctxWithTimeout
 
 	return GetProtocolStateVariables(c.TaikoL1, opts)
 }
diff --git a/pkg/rpc/utils.go b/pkg/rpc/utils.go
index 97e054610..b3b809329 100644
--- a/pkg/rpc/utils.go
+++ b/pkg/rpc/utils.go
@@ -16,8 +16,10 @@ import (
 	"github.com/ethereum/go-ethereum/ethclient/gethclient"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/rpc"
+
 	"github.com/taikoxyz/taiko-client/bindings"
 	"github.com/taikoxyz/taiko-client/bindings/encoding"
+	"github.com/taikoxyz/taiko-client/common/utils"
 )
 
 var (
@@ -377,13 +379,12 @@ func IsArchiveNode(ctx context.Context, client *EthClient, l2GenesisHeight uint6
 // and otherwise returns the context as passed in. cancel func is always set to an empty function
 // so is safe to defer the cancel.
 func ctxWithTimeoutOrDefault(ctx context.Context, defaultTimeout time.Duration) (context.Context, context.CancelFunc) {
-	var (
-		ctxWithTimeout                    = ctx
-		cancel         context.CancelFunc = func() {}
-	)
+	if utils.IsNil(ctx) {
+		return context.WithTimeout(context.Background(), defaultTimeout)
+	}
 	if _, ok := ctx.Deadline(); !ok {
-		ctxWithTimeout, cancel = context.WithTimeout(ctx, defaultTimeout)
+		return context.WithTimeout(ctx, defaultTimeout)
 	}
 
-	return ctxWithTimeout, cancel
+	return ctx, func() {}
 }