[real-time] fix gradual draining of connected clients (#35542)

* [real-time] add regression tests for gradual draining

* [real-time] fix gradual draining of connected clients

GitOrigin-RevId: d9985d2ec72176b567c17f066551cd4f205b0219

services/real-time/app/js/DrainManager.js

@@ -45,7 +45,8 @@ export default {
         client.emit('reconnectGracefully')
         drainedCount++
       }
-      const haveDrainedNClients = drainedCount === N
+      // the rate may be fractional, e.g. rate=27.04 -> reconnect 28 clients
+      const haveDrainedNClients = drainedCount >= N
       if (haveDrainedNClients) {
         break
       }

services/real-time/test/acceptance/js/DrainManagerTests.js

@@ -103,4 +103,63 @@ describe('DrainManagerTests', function () {
       })
     })
   })
+
+  describe('with three clients in the project', function () {
+    beforeEach(function (done) {
+      async.series(
+        [
+          cb => {
+            this.clientA = RealTimeClient.connect(this.project_id, cb)
+          },
+
+          cb => {
+            this.clientB = RealTimeClient.connect(this.project_id, cb)
+          },
+
+          cb => {
+            this.clientC = RealTimeClient.connect(this.project_id, cb)
+          },
+        ],
+        done
+      )
+    })
+
+    describe('starting to drain with a fractional rate', function () {
+      beforeEach(function (done) {
+        this.timeout(10000)
+        this.reconnectedAt = []
+        const trackReconnect = client => cb =>
+          client.on('reconnectGracefully', () => {
+            this.reconnectedAt.push(Date.now())
+            cb()
+          })
+        async.parallel(
+          [
+            trackReconnect(this.clientA),
+            trackReconnect(this.clientB),
+            trackReconnect(this.clientC),
+
+            cb =>
+              drain(1.5)
+                .then(() => cb())
+                .catch(cb),
+          ],
+          done
+        )
+      })
+
+      afterEach(async function () {
+        // reset drain
+        await drain(0)
+      })
+
+      it('should spread the reconnects over multiple polling intervals', function () {
+        // At a rate of 1.5 clients per second, two clients should be asked to
+        // reconnect in the first 1s polling interval and the third client one
+        // polling interval later.
+        const sorted = this.reconnectedAt.slice().sort((a, b) => a - b)
+        expect(sorted[2] - sorted[0]).to.be.greaterThan(500)
+      })
+    })
+  })
 })

services/real-time/test/unit/js/DrainManager.test.js

@@ -43,6 +43,14 @@ describe('DrainManager', function () {
         resolve()
       })
     })
+
+    it('should set a fractional drain rate when clients do not divide evenly', function (ctx) {
+      ctx.clients.push({ id: 5400, emit: sinon.stub() })
+      ctx.DrainManager.startDrainTimeWindow(ctx.io, 9)
+      ctx.DrainManager.startDrain
+        .calledWith(ctx.io, 10.001851851851852)
+        .should.equal(true)
+    })
   })
 
   describe('reconnectNClients', function () {
@@ -122,5 +130,92 @@ describe('DrainManager', function () {
         })
       })
     })
+
+    describe('with a fractional rate', function () {
+      describe('after first pass', function () {
+        beforeEach(function (ctx) {
+          ctx.requestedAllClientsToReconnect =
+            ctx.DrainManager.reconnectNClients(ctx.io, 3.5)
+        })
+
+        it('should reconnect the first 4 clients', function (ctx) {
+          ;[0, 1, 2, 3].map(i =>
+            ctx.clients[i].emit
+              .calledWith('reconnectGracefully')
+              .should.equal(true)
+          )
+        })
+
+        it('should not reconnect any more clients', function (ctx) {
+          ;[4, 5, 6, 7, 8, 9].map(i =>
+            ctx.clients[i].emit
+              .calledWith('reconnectGracefully')
+              .should.equal(false)
+          )
+        })
+
+        it('should not report that all clients have been told to reconnect', function (ctx) {
+          ctx.requestedAllClientsToReconnect.should.equal(false)
+        })
+
+        describe('after second pass', function () {
+          beforeEach(function (ctx) {
+            ctx.requestedAllClientsToReconnect =
+              ctx.DrainManager.reconnectNClients(ctx.io, 3.5)
+          })
+
+          it('should reconnect the next 4 clients', function (ctx) {
+            ;[4, 5, 6, 7].map(i =>
+              ctx.clients[i].emit
+                .calledWith('reconnectGracefully')
+                .should.equal(true)
+            )
+          })
+
+          it('should not reconnect any more clients', function (ctx) {
+            ;[8, 9].map(i =>
+              ctx.clients[i].emit
+                .calledWith('reconnectGracefully')
+                .should.equal(false)
+            )
+          })
+
+          it('should not reconnect the first 4 clients again', function (ctx) {
+            ;[0, 1, 2, 3].map(i =>
+              ctx.clients[i].emit.calledOnce.should.equal(true)
+            )
+          })
+
+          it('should not report that all clients have been told to reconnect', function (ctx) {
+            ctx.requestedAllClientsToReconnect.should.equal(false)
+          })
+
+          describe('after third pass', function () {
+            beforeEach(function (ctx) {
+              ctx.requestedAllClientsToReconnect =
+                ctx.DrainManager.reconnectNClients(ctx.io, 3.5)
+            })
+
+            it('should reconnect the last 2 clients', function (ctx) {
+              ;[8, 9].map(i =>
+                ctx.clients[i].emit
+                  .calledWith('reconnectGracefully')
+                  .should.equal(true)
+              )
+            })
+
+            it('should not reconnect the first 8 clients again', function (ctx) {
+              ;[0, 1, 2, 3, 4, 5, 6, 7].map(i =>
+                ctx.clients[i].emit.calledOnce.should.equal(true)
+              )
+            })
+
+            it('should report that all clients have been told to reconnect', function (ctx) {
+              ctx.requestedAllClientsToReconnect.should.equal(true)
+            })
+          })
+        })
+      })
+    })
   })
 })