From c754b206c21ee2d374203bc83f723123a476fe24 Mon Sep 17 00:00:00 2001
From: testinginprod <98415576+testinginprod@users.noreply.github.com>
Date: Wed, 30 Oct 2024 18:23:34 +0100
Subject: [PATCH] refactor(serverv2): remove unused interface methods, honuor
 context  (#22394)

---
 runtime/v2/store.go          | 23 -----------------------
 server/v2/api/grpc/server.go |  8 ++++----
 server/v2/api/utils.go       | 22 ++++++++++++++++++++++
 server/v2/api/utils_test.go  | 35 +++++++++++++++++++++++++++++++++++
 4 files changed, 61 insertions(+), 27 deletions(-)
 create mode 100644 server/v2/api/utils.go
 create mode 100644 server/v2/api/utils_test.go

diff --git a/runtime/v2/store.go b/runtime/v2/store.go
index 5d37c321c1fe..0d9b4422844f 100644
--- a/runtime/v2/store.go
+++ b/runtime/v2/store.go
@@ -7,7 +7,6 @@ import (
 	"cosmossdk.io/core/store"
 	"cosmossdk.io/server/v2/stf"
 	storev2 "cosmossdk.io/store/v2"
-	"cosmossdk.io/store/v2/proof"
 )
 
 // NewKVStoreService creates a new KVStoreService.
@@ -29,34 +28,12 @@ type Store interface {
 	// version. Must error when the version does not exist.
 	StateAt(version uint64) (store.ReaderMap, error)
 
-	// SetInitialVersion sets the initial version of the store.
-	SetInitialVersion(uint64) error
-
-	// WorkingHash writes the provided changeset to the state and returns
-	// the working hash of the state.
-	WorkingHash(changeset *store.Changeset) (store.Hash, error)
-
-	// Commit commits the provided changeset and returns the new state root of the state.
-	Commit(changeset *store.Changeset) (store.Hash, error)
-
-	// Query is a key/value query directly to the underlying database. This skips the appmanager
-	Query(storeKey []byte, version uint64, key []byte, prove bool) (storev2.QueryResult, error)
-
-	// GetStateStorage returns the SS backend.
-	GetStateStorage() storev2.VersionedWriter
-
-	// GetStateCommitment returns the SC backend.
-	GetStateCommitment() storev2.Committer
-
 	// LoadVersion loads the RootStore to the given version.
 	LoadVersion(version uint64) error
 
 	// LoadLatestVersion behaves identically to LoadVersion except it loads the
 	// latest version implicitly.
 	LoadLatestVersion() error
-
-	// LastCommitID returns the latest commit ID
-	LastCommitID() (proof.CommitID, error)
 }
 
 // StoreLoader allows for custom loading of the store, this is useful when upgrading the store from a previous version
diff --git a/server/v2/api/grpc/server.go b/server/v2/api/grpc/server.go
index 4768b74ca1c8..e163bf3183c1 100644
--- a/server/v2/api/grpc/server.go
+++ b/server/v2/api/grpc/server.go
@@ -25,6 +25,7 @@ import (
 	"cosmossdk.io/core/transaction"
 	"cosmossdk.io/log"
 	serverv2 "cosmossdk.io/server/v2"
+	"cosmossdk.io/server/v2/api"
 	"cosmossdk.io/server/v2/api/grpc/gogoreflection"
 )
 
@@ -197,13 +198,13 @@ func (s *Server[T]) Config() any {
 	return s.config
 }
 
-func (s *Server[T]) Start(context.Context) error {
+func (s *Server[T]) Start(ctx context.Context) error {
 	if !s.config.Enable {
 		s.logger.Info(fmt.Sprintf("%s server is disabled via config", s.Name()))
 		return nil
 	}
 
-	listener, err := net.Listen("tcp", s.config.Address)
+	listener, err := (&net.ListenConfig{}).Listen(ctx, "tcp", s.config.Address)
 	if err != nil {
 		return fmt.Errorf("failed to listen on address %s: %w", s.config.Address, err)
 	}
@@ -222,8 +223,7 @@ func (s *Server[T]) Stop(ctx context.Context) error {
 	}
 
 	s.logger.Info("stopping gRPC server...", "address", s.config.Address)
-	s.grpcSrv.GracefulStop()
-	return nil
+	return api.DoUntilCtxExpired(ctx, s.grpcSrv.GracefulStop)
 }
 
 // GetGRPCServer returns the underlying gRPC server.
diff --git a/server/v2/api/utils.go b/server/v2/api/utils.go
new file mode 100644
index 000000000000..a6bef1224db5
--- /dev/null
+++ b/server/v2/api/utils.go
@@ -0,0 +1,22 @@
+package api
+
+import "context"
+
+// DoUntilCtxExpired runs the given function until the context is expired or
+// the function exits.
+// This forces context to be honored.
+func DoUntilCtxExpired(ctx context.Context, f func()) error {
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+
+		f()
+	}()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-done:
+		return nil
+	}
+}
diff --git a/server/v2/api/utils_test.go b/server/v2/api/utils_test.go
new file mode 100644
index 000000000000..3c181d7653d1
--- /dev/null
+++ b/server/v2/api/utils_test.go
@@ -0,0 +1,35 @@
+package api
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestDoUntilCtxExpired(t *testing.T) {
+	t.Run("success", func(t *testing.T) {
+		ctx := context.Background()
+
+		funcRan := false
+		err := DoUntilCtxExpired(ctx, func() {
+			funcRan = true
+		})
+		require.NoError(t, err)
+		require.True(t, funcRan)
+	})
+
+	t.Run("context expired", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+
+		funcRan := false
+		err := DoUntilCtxExpired(ctx, func() {
+			cancel()
+			funcRan = true
+			<-time.After(time.Second)
+		})
+		require.ErrorIs(t, err, context.Canceled)
+		require.True(t, funcRan)
+	})
+}