From 775c56e87e742a53f5d117601346c691ac2a4b08 Mon Sep 17 00:00:00 2001
From: Ekaterina Pavlova <ekt@morphbits.io>
Date: Sun, 18 Feb 2024 15:27:52 +0300
Subject: [PATCH] network: ensure server is started and shut down only once

Use started atomic.Bool field to ensure that the node server shutdown
procedure is executed only once. Prevent the following panic caused by
server double-shutdown in testing code:
```
--- FAIL: TestServerRegisterPeer (0
.06s)
 panic: closed twice
 goroutine 60 [running]:
 testing.tRunner.func1.2({0x104c40b20, 0x104d0ec90})
 	/opt/homebrew/opt/go/libexec/src/testing/testing.go:1545 +0x1c8
 testing.tRunner.func1()
 	/opt/homebrew/opt/go/libexec/src/testing/testing.go:1548 +0x360
 panic({0x104c40b20?, 0x104d0ec90?})
 	/opt/homebrew/opt/go/libexec/src/runtime/panic.go:914 +0x218
 github.com/nspcc-dev/neo-go/pkg/network.(*fakeTransp).Close
 (0x14000159e08?)
 	/Users/ekaterinapavlova/Workplace/neo-go/pkg/network
 	/discovery_test.go:83 +0x54
 github.com/nspcc-dev/neo-go/pkg/network.(*Server).Shutdown
 (0x14000343400)
 	/Users/ekaterinapavlova/Workplace/neo-go/pkg/network/server.go:299
 	 +0x104
 github.com/nspcc-dev/neo-go/pkg/network.startWithCleanup.func1()
 	/Users/ekaterinapavlova/Workplace/neo-go/pkg/network/server_test
 	.go:408 +0x20
 testing.(*common).Cleanup.func1()
 	/opt/homebrew/opt/go/libexec/src/testing/testing.go:1169 +0x110
 testing.(*common).runCleanup(0x1400032c340, 0x14000159d80?)
 	/opt/homebrew/opt/go/libexec/src/testing/testing.go:1347 +0xd8
 testing.tRunner.func2()
 	/opt/homebrew/opt/go/libexec/src/testing/testing.go:1589 +0x2c
 testing.tRunner(0x1400032c340, 0x104d0c5d0)
 	/opt/homebrew/opt/go/libexec/src/testing/testing.go:1601 +0x114
 created by testing.(*T).Run in goroutine 1
 	/opt/homebrew/opt/go/libexec/src/testing/testing.go:1648 +0x33c
```

Signed-off-by: Ekaterina Pavlova <ekt@morphbits.io>
---
 pkg/network/server.go      | 14 ++++++++++++--
 pkg/network/server_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/pkg/network/server.go b/pkg/network/server.go
index 0e919cbdd..b97cf2dd1 100644
--- a/pkg/network/server.go
+++ b/pkg/network/server.go
@@ -138,6 +138,9 @@ type (
 		stateSync StateSync
 
 		log *zap.Logger
+
+		// started used to Start and Shutdown server only once.
+		started atomic.Bool
 	}
 
 	peerDrop struct {
@@ -262,8 +265,12 @@ func (s *Server) ID() uint32 {
 }
 
 // Start will start the server and its underlying transport. Calling it twice
-// is an error. Caller should wait for Start to finish for normal server operation.
+// is a no-op. Caller should wait for Start to finish for normal server operation.
 func (s *Server) Start() {
+	if !s.started.CompareAndSwap(false, true) {
+		s.log.Info("node server already started")
+		return
+	}
 	s.log.Info("node started",
 		zap.Uint32("blockHeight", s.chain.BlockHeight()),
 		zap.Uint32("headerHeight", s.chain.HeaderHeight()))
@@ -288,9 +295,12 @@ func (s *Server) Start() {
 	go s.run()
 }
 
-// Shutdown disconnects all peers and stops listening. Calling it twice is an error,
+// Shutdown disconnects all peers and stops listening. Calling it twice is a no-op,
 // once stopped the same intance of the Server can't be started again by calling Start.
 func (s *Server) Shutdown() {
+	if !s.started.CompareAndSwap(true, false) {
+		return
+	}
 	s.log.Info("shutting down server", zap.Int("peers", s.PeerCount()))
 	for _, tr := range s.transports {
 		tr.Close()
diff --git a/pkg/network/server_test.go b/pkg/network/server_test.go
index 7dbce881d..c1b5d96dc 100644
--- a/pkg/network/server_test.go
+++ b/pkg/network/server_test.go
@@ -96,10 +96,12 @@ func TestServerStartAndShutdown(t *testing.T) {
 		require.Eventually(t, func() bool { return 1 == s.PeerCount() }, time.Second, time.Millisecond*10)
 
 		assert.True(t, s.transports[0].(*fakeTransp).started.Load())
+		require.True(t, s.started.Load())
 		assert.Nil(t, s.txCallback)
 
 		s.Shutdown()
 
+		require.False(t, s.started.Load())
 		require.True(t, s.transports[0].(*fakeTransp).closed.Load())
 		err, ok := p.droppedWith.Load().(error)
 		require.True(t, ok)
@@ -115,11 +117,34 @@ func TestServerStartAndShutdown(t *testing.T) {
 		s.register <- p
 
 		assert.True(t, s.services["fake"].(*fakeConsensus).started.Load())
+		require.True(t, s.started.Load())
 
 		s.Shutdown()
 
+		require.False(t, s.started.Load())
 		require.True(t, s.services["fake"].(*fakeConsensus).stopped.Load())
 	})
+	t.Run("double start", func(t *testing.T) {
+		s := newTestServer(t, ServerConfig{})
+		startWithCleanup(t, s)
+
+		// Attempt to start the server again.
+		s.Start()
+
+		require.True(t, s.started.Load(), "server should still be marked as started after second Start call")
+	})
+	t.Run("double shutdown", func(t *testing.T) {
+		s := newTestServer(t, ServerConfig{})
+		s.Start()
+		require.True(t, s.started.Load(), "server should still be marked as started after second Start call")
+		s.Shutdown()
+
+		require.False(t, s.started.Load(), "server should be marked as not started after second Shutdown call")
+		// Attempt to shutdown the server again.
+		s.Shutdown()
+		// Verify the server state remains unchanged and is still considered shutdown.
+		require.False(t, s.started.Load(), "server should remain shutdown after second call")
+	})
 }
 
 func TestServerRegisterPeer(t *testing.T) {