Repository: nats-io/nats-server (go)

Issue:
Client connections load balancing among nats cluster nodes
Hello! Currently the mechanism of load balancing of client connections to nats cluster is based on random node selection. The problem (already discussed in some of the issues, like https://github.com/nats-io/nats-server/issues/1359) is when we restart nats cluster. To be able to serve clients we restart cluster one node after another and this will lead us to approximately the following distribution (step-by-step example for 3 node cluster):
1) original connection distribution is 24 / 24 / 24
2) after node1 restart : 0 / 36 / 36
3) after node2 restart:  18 / 0 / 54
4) after node3 restart: 45 / 27 / 0

As you can see original uniform distribution 1/3 | 1/3 | 1/3 tends to 2/3 | 1/3 | 0 . I suggest you to add some modifications that will allow developers to get around this problem (for example, we use nats cluster in 24/7 fashion and never stop it completely but sometimes need to upgrade server version or our hardware). There can be several solutions:
1) Implement clever load balancing (when nats cluster node tells client to try to connect to another known node which have less active clients)
2) Provide a command which can be sent to cluster and cause it to disconnect existing clients - we can use it after cluster restart is finished to redistribute client connections among nodes.

Pull request:
[ADDED] $SYS server request to 'kick' or 'LDM' a client connection
 - [X] Link to issue, e.g. `Resolves #NNN`
 - [X] Branch rebased on top of current main (`git pull --rebase origin main`)
 - [ ] Changes squashed to a single commit (described [here](http://gitready.com/advanced/2009/02/10/squashing-commits-with-rebase.html))
 - [x] Build is green in Travis CI
 - [X] You have certified that the contribution is your original work and that you license the work to the project under the [Apache 2 license](https://github.com/nats-io/nats-server/blob/main/LICENSE)

Resolves #1556

### Changes proposed in this pull request:

Adds tw new $SYS server API endpoints:

- `$SYS.REQ.SERVER.%s.KICK` (where %s is the server_id) which 'kicks' (effectiveley 'rebalance' as the client application reconnects itself right away (potentially to another server in the cluster)). The service takes a JSON payload containing either an "id" or a "name" field. "id" disconnects the client connection id, "name" disconnects _all_ of the clients connected to the server with that name.

- `$SYS.REQ.SERVER.%s.LDM` (where %s is the server_id) and takes a JSON payload containing either an "id" or a "name" field. "id" sends an LDM Info message to the client connection id, "name" sends an LDM Info message to _all_ of the clients connected to the server with that name.

This features allow administrators to manually 're-balance' client connections between the servers in the cluster (e.g. after a rolling upgrade of the servers where one server ends up with no client connections after the upgrade), by kicking some of the client connections from one of the 'overloaded' (in comparison to other servers) servers in the cluster, causing them to re-estalibsh their connection to (hopefully) another server.

Patch:
diff --git a/server/client.go b/server/client.go
index c0f9e844..41997db5 100644
--- a/server/client.go
+++ b/server/client.go
@@ -213,6 +213,7 @@ const (
 	DuplicateServerName
 	MinimumVersionRequired
 	ClusterNamesIdentical
+	Kicked
 )
 
 // Some flags passed to processMsgResults
diff --git a/server/events.go b/server/events.go
index 75cd9aa0..58437f70 100644
--- a/server/events.go
+++ b/server/events.go
@@ -57,6 +57,8 @@ const (
 	accConnsEventSubjNew      = "$SYS.ACCOUNT.%s.SERVER.CONNS"
 	accConnsEventSubjOld      = "$SYS.SERVER.ACCOUNT.%s.CONNS" // kept for backward compatibility
 	shutdownEventSubj         = "$SYS.SERVER.%s.SHUTDOWN"
+	clientKickReqSubj         = "$SYS.REQ.SERVER.%s.KICK"
+	clientLDMReqSubj          = "$SYS.REQ.SERVER.%s.LDM"
 	authErrorEventSubj        = "$SYS.SERVER.%s.CLIENT.AUTH.ERR"
 	authErrorAccountEventSubj = "$SYS.ACCOUNT.CLIENT.AUTH.ERR"
 	serverStatsSubj           = "$SYS.SERVER.%s.STATSZ"
@@ -1228,6 +1230,17 @@ func (s *Server) initEventTracking() {
 	if _, err := s.sysSubscribe(subject, s.noInlineCallback(s.reloadConfig)); err != nil {
 		s.Errorf("Error setting up server reload handler: %v", err)
 	}
+
+	// Client connection kick
+	subject = fmt.Sprintf(clientKickReqSubj, s.info.ID)
+	if _, err := s.sysSubscribe(subject, s.noInlineCallback(s.kickClient)); err != nil {
+		s.Errorf("Error setting up client kick service: %v", err)
+	}
+	// Client connection LDM
+	subject = fmt.Sprintf(clientLDMReqSubj, s.info.ID)
+	if _, err := s.sysSubscribe(subject, s.noInlineCallback(s.ldmClient)); err != nil {
+		s.Errorf("Error setting up client LDM service: %v", err)
+	}
 }
 
 // UserInfo returns basic information to a user about bound account and user permissions.
@@ -2714,6 +2727,49 @@ func (s *Server) reloadConfig(sub *subscription, c *client, _ *Account, subject,
 	})
 }
 
+type KickClientReq struct {
+	CID uint64 `json:"cid"`
+}
+
+type LDMClientReq struct {
+	CID uint64 `json:"cid"`
+}
+
+func (s *Server) kickClient(_ *subscription, c *client, _ *Account, subject, reply string, hdr, msg []byte) {
+	if !s.eventsRunning() {
+		return
+	}
+
+	var req KickClientReq
+	if err := json.Unmarshal(msg, &req); err != nil {
+		s.sys.client.Errorf("Error unmarshalling kick client request: %v", err)
+		return
+	}
+
+	optz := &EventFilterOptions{}
+	s.zReq(c, reply, hdr, msg, optz, optz, func() (interface{}, error) {
+		return nil, s.DisconnectClientByID(req.CID)
+	})
+
+}
+
+func (s *Server) ldmClient(_ *subscription, c *client, _ *Account, subject, reply string, hdr, msg []byte) {
+	if !s.eventsRunning() {
+		return
+	}
+
+	var req LDMClientReq
+	if err := json.Unmarshal(msg, &req); err != nil {
+		s.sys.client.Errorf("Error unmarshalling kick client request: %v", err)
+		return
+	}
+
+	optz := &EventFilterOptions{}
+	s.zReq(c, reply, hdr, msg, optz, optz, func() (interface{}, error) {
+		return nil, s.LDMClientByID(req.CID)
+	})
+}
+
 // Helper to grab account name for a client.
 func accForClient(c *client) string {
 	if c.acc != nil {
diff --git a/server/monitor.go b/server/monitor.go
index 7721fd6d..2579fe1d 100644
--- a/server/monitor.go
+++ b/server/monitor.go
@@ -2416,6 +2416,8 @@ func (reason ClosedState) String() string {
 		return "Minimum Version Required"
 	case ClusterNamesIdentical:
 		return "Cluster Names Identical"
+	case Kicked:
+		return "Kicked"
 	}
 
 	return "Unknown State"
diff --git a/server/server.go b/server/server.go
index b677eab7..c858974d 100644
--- a/server/server.go
+++ b/server/server.go
@@ -4252,3 +4252,36 @@ func (s *Server) changeRateLimitLogInterval(d time.Duration) {
 	default:
 	}
 }
+
+// DisconnectClientByID disconnects a client by connection ID
+func (s *Server) DisconnectClientByID(id uint64) error {
+	client := s.clients[id]
+	if client != nil {
+		client.closeConnection(Kicked)
+		return nil
+	}
+	return errors.New("no such client id")
+}
+
+// LDMClientByID sends a Lame Duck Mode info message to a client by connection ID
+func (s *Server) LDMClientByID(id uint64) error {
+	info := s.copyInfo()
+	info.LameDuckMode = true
+
+	c := s.clients[id]
+	if c != nil {
+		c.mu.Lock()
+		defer c.mu.Unlock()
+		if c.opts.Protocol >= ClientProtoInfo &&
+			c.flags.isSet(firstPongSent) {
+			// sendInfo takes care of checking if the connection is still
+			// valid or not, so don't duplicate tests here.
+			c.Debugf("sending Lame Duck Mode info to client")
+			c.enqueueProto(c.generateClientInfoJSON(info))
+			return nil
+		} else {
+			return errors.New("ClientProtoInfo < ClientOps.Protocol or first pong not sent")
+		}
+	}
+	return errors.New("no such client id")
+}


Test patch:
diff --git a/server/events_test.go b/server/events_test.go
index a4cb9555..e3541999 100644
--- a/server/events_test.go
+++ b/server/events_test.go
@@ -1668,7 +1668,7 @@ func TestSystemAccountWithGateways(t *testing.T) {
 
 	// If this tests fails with wrong number after 10 seconds we may have
 	// added a new inititial subscription for the eventing system.
-	checkExpectedSubs(t, 53, sa)
+	checkExpectedSubs(t, 55, sa)
 
 	// Create a client on B and see if we receive the event
 	urlb := fmt.Sprintf("nats://%s:%d", ob.Host, ob.Port)
@@ -3412,6 +3412,66 @@ func TestServerEventsReload(t *testing.T) {
 	require_True(t, s.getOpts().PingInterval == 200*time.Millisecond)
 }
 
+func TestServerEventsLDMKick(t *testing.T) {
+	ldmed := make(chan bool, 1)
+	disconnected := make(chan bool, 1)
+
+	s, opts := runTrustedServer(t)
+	defer s.Shutdown()
+
+	acc, akp := createAccount(s)
+	s.setSystemAccount(acc)
+
+	url := fmt.Sprintf("nats://%s:%d", opts.Host, opts.Port)
+	ncs, err := nats.Connect(url, createUserCreds(t, s, akp))
+	if err != nil {
+		t.Fatalf("Error on connect: %v", err)
+	}
+	defer ncs.Close()
+
+	_, akp2 := createAccount(s)
+
+	nc, err := nats.Connect(url, createUserCreds(t, s, akp2), nats.Name("TEST EVENTS LDM+KICK"), nats.LameDuckModeHandler(func(_ *nats.Conn) { ldmed <- true }))
+	if err != nil {
+		t.Fatalf("Error on connect: %v", err)
+	}
+	defer nc.Close()
+
+	nc.SetDisconnectErrHandler(func(_ *nats.Conn, err error) { disconnected <- true })
+
+	cid, err := nc.GetClientID()
+	if err != nil {
+		t.Fatalf("Error on getting the CID: %v", err)
+	}
+
+	reqldm := LDMClientReq{CID: cid}
+	reqldmpayload, _ := json.Marshal(reqldm)
+	reqkick := KickClientReq{CID: cid}
+	reqkickpayload, _ := json.Marshal(reqkick)
+
+	_, err = ncs.Request(fmt.Sprintf("$SYS.REQ.SERVER.%s.LDM", s.ID()), reqldmpayload, time.Second)
+	if err != nil {
+		t.Fatalf("Error trying to publish the LDM request: %v", err)
+	}
+
+	select {
+	case <-ldmed:
+	case <-time.After(time.Second):
+		t.Fatalf("timeout waiting for the connection to receive the LDM signal")
+	}
+
+	_, err = ncs.Request(fmt.Sprintf("$SYS.REQ.SERVER.%s.KICK", s.ID()), reqkickpayload, time.Second)
+	if err != nil {
+		t.Fatalf("Error trying to publish the KICK request: %v", err)
+	}
+
+	select {
+	case <-disconnected:
+	case <-time.After(time.Second):
+		t.Fatalf("timeout waiting for the client to get disconnected")
+	}
+}
+
 func Benchmark_GetHash(b *testing.B) {
 	b.StopTimer()
 	// Get 100 random names
diff --git a/server/monitor_test.go b/server/monitor_test.go
index 9d9f7a2d..d9155dd4 100644
--- a/server/monitor_test.go
+++ b/server/monitor_test.go
@@ -3946,7 +3946,7 @@ func TestMonitorAccountz(t *testing.T) {
 	body = string(readBody(t, fmt.Sprintf("http://127.0.0.1:%d%s?acc=$SYS", s.MonitorAddr().Port, AccountzPath)))
 	require_Contains(t, body, `"account_detail": {`)
 	require_Contains(t, body, `"account_name": "$SYS",`)
-	require_Contains(t, body, `"subscriptions": 47,`)
+	require_Contains(t, body, `"subscriptions": 49,`)
 	require_Contains(t, body, `"is_system": true,`)
 	require_Contains(t, body, `"system_account": "$SYS"`)