From ffda2c64c4260e1dd2cc2760faf1f8d33bdcf670 Mon Sep 17 00:00:00 2001
From: s7v7nislands <s7v7nislands@gmail.com>
Date: Thu, 4 May 2023 16:54:45 +0800
Subject: [PATCH] rpc: use atomic types (#27214)

rpc: use atomic type
---
 rpc/client.go          |  4 ++--
 rpc/server.go          | 10 +++++-----
 signer/core/stdioui.go |  4 ++--
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/rpc/client.go b/rpc/client.go
index 69ff4851e..fbeb5a181 100644
--- a/rpc/client.go
+++ b/rpc/client.go
@@ -79,7 +79,7 @@ type Client struct {
 	isHTTP   bool      // connection type: http, ws or ipc
 	services *serviceRegistry
 
-	idCounter uint32
+	idCounter atomic.Uint32
 
 	// This function, if non-nil, is called when the connection is lost.
 	reconnectFunc reconnectFunc
@@ -263,7 +263,7 @@ func (c *Client) RegisterName(name string, receiver interface{}) error {
 }
 
 func (c *Client) nextID() json.RawMessage {
-	id := atomic.AddUint32(&c.idCounter, 1)
+	id := c.idCounter.Add(1)
 	return strconv.AppendUint(nil, uint64(id), 10)
 }
 
diff --git a/rpc/server.go b/rpc/server.go
index 9c72c26d7..089bbb1fd 100644
--- a/rpc/server.go
+++ b/rpc/server.go
@@ -48,7 +48,7 @@ type Server struct {
 
 	mutex  sync.Mutex
 	codecs map[ServerCodec]struct{}
-	run    int32
+	run    atomic.Bool
 }
 
 // NewServer creates a new server instance with no registered handlers.
@@ -56,8 +56,8 @@ func NewServer() *Server {
 	server := &Server{
 		idgen:  randomIDGenerator(),
 		codecs: make(map[ServerCodec]struct{}),
-		run:    1,
 	}
+	server.run.Store(true)
 	// Register the default service providing meta information about the RPC service such
 	// as the services and methods it offers.
 	rpcService := &RPCService{server}
@@ -95,7 +95,7 @@ func (s *Server) trackCodec(codec ServerCodec) bool {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
-	if atomic.LoadInt32(&s.run) == 0 {
+	if !s.run.Load() {
 		return false // Don't serve if server is stopped.
 	}
 	s.codecs[codec] = struct{}{}
@@ -114,7 +114,7 @@ func (s *Server) untrackCodec(codec ServerCodec) {
 // this mode.
 func (s *Server) serveSingleRequest(ctx context.Context, codec ServerCodec) {
 	// Don't serve if server is stopped.
-	if atomic.LoadInt32(&s.run) == 0 {
+	if !s.run.Load() {
 		return
 	}
 
@@ -144,7 +144,7 @@ func (s *Server) Stop() {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
-	if atomic.CompareAndSwapInt32(&s.run, 1, 0) {
+	if s.run.CompareAndSwap(true, false) {
 		log.Debug("RPC server shutting down")
 		for codec := range s.codecs {
 			codec.close()
diff --git a/signer/core/stdioui.go b/signer/core/stdioui.go
index 6963a8912..a0ce68441 100644
--- a/signer/core/stdioui.go
+++ b/signer/core/stdioui.go
@@ -25,7 +25,7 @@ import (
 )
 
 type StdIOUI struct {
-	client rpc.Client
+	client *rpc.Client
 }
 
 func NewStdIOUI() *StdIOUI {
@@ -33,7 +33,7 @@ func NewStdIOUI() *StdIOUI {
 	if err != nil {
 		log.Crit("Could not create stdio client", "err", err)
 	}
-	ui := &StdIOUI{client: *client}
+	ui := &StdIOUI{client: client}
 	return ui
 }