Reject new sessions if shutdown is scheduled, resume will still work.

hub.go

@@ -62,6 +62,7 @@ var (
 	NoSuchSession       = NewError("no_such_session", "The session to resume does not exist.")
 	TokenNotValidYet    = NewError("token_not_valid_yet", "The token is not valid yet.")
 	TokenExpired        = NewError("token_expired", "The token is expired.")
+	ShutdownScheduled   = NewError("shutdown_scheduled", "The server is shutting down.")
 
 	// Maximum number of concurrent requests to a backend.
 	defaultMaxConcurrentRequestsPerHost = 8
@@ -838,6 +839,12 @@ func (h *Hub) processRegister(client *Client, message *ClientMessage, backend *B
 		return
 	}
 
+	if h.shutdownScheduled.Load() {
+		session.Close()
+		client.SendMessage(message.NewErrorServerMessage(ShutdownScheduled))
+		return
+	}
+
 	session.SetClient(client)
 	h.sessions[sessionIdData.Sid] = session
 	h.clients[sessionIdData.Sid] = client
@@ -1014,6 +1021,11 @@ func (h *Hub) processHello(client *Client, message *ClientMessage) {
 		return
 	}
 
+	if h.shutdownScheduled.Load() {
+		client.SendMessage(message.NewErrorServerMessage(ShutdownScheduled))
+		return
+	}
+
 	// Make sure client doesn't get disconnected while calling auth backend.
 	h.mu.Lock()
 	delete(h.expectHelloClients, client)

hub_test.go

@@ -5613,7 +5613,8 @@ func TestGracefulShutdownOnExpiration(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	if _, err := client.RunUntilHello(ctx); err != nil {
+	hello, err := client.RunUntilHello(ctx)
+	if err != nil {
 		t.Error(err)
 	}
 
@@ -5625,12 +5626,63 @@ func TestGracefulShutdownOnExpiration(t *testing.T) {
 	}
 
 	client.Close()
+	if err := client.WaitForClientRemoved(ctx); err != nil {
+		t.Error(err)
+	}
+
 	select {
 	case <-hub.ShutdownChannel():
 		t.Error("should not have shutdown")
 	case <-time.After(100 * time.Millisecond):
 	}
 
+	client = NewTestClient(t, server, hub)
+	defer client.CloseWithBye()
+
+	if err := client.SendHelloResume(hello.Hello.ResumeId); err != nil {
+		t.Fatal(err)
+	}
+	hello2, err := client.RunUntilHello(ctx)
+	if err != nil {
+		t.Error(err)
+	} else {
+		if hello2.Hello.UserId != testDefaultUserId {
+			t.Errorf("Expected \"%s\", got %+v", testDefaultUserId, hello2.Hello)
+		}
+		if hello2.Hello.SessionId != hello.Hello.SessionId {
+			t.Errorf("Expected session id %s, got %+v", hello.Hello.SessionId, hello2.Hello)
+		}
+		if hello2.Hello.ResumeId != hello.Hello.ResumeId {
+			t.Errorf("Expected resume id %s, got %+v", hello.Hello.ResumeId, hello2.Hello)
+		}
+	}
+
+	client.Close()
+	if err := client.WaitForClientRemoved(ctx); err != nil {
+		t.Error(err)
+	}
+
+	select {
+	case <-hub.ShutdownChannel():
+		t.Error("should not have shutdown")
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	client2 := NewTestClient(t, server, hub)
+	defer client2.CloseWithBye()
+
+	if err := client2.SendHello(testDefaultUserId + "2"); err != nil {
+		t.Fatal(err)
+	}
+
+	if msg, err := client2.RunUntilMessage(ctx); err != nil {
+		t.Fatal(err)
+	} else if msg.Type != "error" {
+		t.Errorf("expected error message, got %+v", msg)
+	} else if msg.Error.Code != "shutdown_scheduled" {
+		t.Errorf("expected error \"shutdown_scheduled\", got %+v", msg)
+	}
+
 	performHousekeeping(hub, time.Now().Add(sessionExpireDuration+time.Second))
 
 	select {