Merge pull request #1137 from interval/begin-shutdown

Add async `safelyClose` method to close after transactions complete

src/classes/IntervalClient.ts

@@ -111,12 +111,14 @@ export default class IntervalClient {
   #httpEndpoint: string
   #logger: Logger
   #completeHttpRequestDelayMs: number = 3000
+  #completeShutdownDelayMs: number = 3000
   #retryIntervalMs: number = 3000
   #pingIntervalMs: number = 30_000
   #closeUnresponsiveConnectionTimeoutMs: number = 3 * 60 * 1000 // 3 minutes
   #reinitializeBatchTimeoutMs: number = 200
   #pingIntervalHandle: NodeJS.Timeout | undefined
   #intentionallyClosed = false
+  #resolveShutdown: (() => void) | undefined
   #config: InternalConfig
 
   #actionDefinitions: ActionDefinition[] = []
@@ -372,7 +374,8 @@ export default class IntervalClient {
     return await result
   }
 
-  close() {
+  immediatelyClose() {
+    this.#resolveShutdown = undefined
     this.#intentionallyClosed = true
 
     if (this.#serverRpc) {
@@ -393,6 +396,36 @@ export default class IntervalClient {
     this.#isConnected = false
   }
 
+  async safelyClose(): Promise<void> {
+    const response = await this.#send(
+      'BEGIN_HOST_SHUTDOWN',
+      {},
+      {
+        attemptPeerSend: false,
+      }
+    )
+
+    if (response.type === 'error') {
+      throw new IntervalError(
+        response.message ?? 'Unknown error sending shutdown request.'
+      )
+    }
+
+    if (this.#ioResponseHandlers.size === 0) {
+      this.immediatelyClose()
+      return
+    }
+
+    return new Promise<void>(resolve => {
+      this.#resolveShutdown = resolve
+    }).then(() => {
+      // doing this here and in #close just to be extra sure
+      // it's not missed in any future code paths
+      this.#resolveShutdown = undefined
+      this.immediatelyClose()
+    })
+  }
+
   async declareHost(httpHostId: string) {
     await this.#walkRoutes()
 
@@ -655,6 +688,12 @@ export default class IntervalClient {
       this.#peerIdMap.delete(transactionId)
       this.#peerIdToTransactionIdsMap.get(peerId)?.delete(transactionId)
     }
+
+    if (this.#resolveShutdown && this.#ioResponseHandlers.size === 0) {
+      setTimeout(() => {
+        this.#resolveShutdown?.()
+      }, this.#completeShutdownDelayMs)
+    }
   }
 
   /**
@@ -756,7 +795,9 @@ export default class IntervalClient {
           new Date().getTime() - this.#closeUnresponsiveConnectionTimeoutMs
         ) {
           this.#logger.warn(
-            'No pong received in last three minutes, closing connection to Interval and retrying...'
+            `No pong received in last ${
+              this.#closeUnresponsiveConnectionTimeoutMs
+            }ms, closing connection to Interval and retrying...`
           )
           if (this.#pingIntervalHandle) {
             clearInterval(this.#pingIntervalHandle)
@@ -917,6 +958,13 @@ export default class IntervalClient {
         }
       },
       START_TRANSACTION: async inputs => {
+        if (this.#resolveShutdown) {
+          this.#logger.debug(
+            'In process of closing, refusing to start transaction'
+          )
+          return
+        }
+
         if (!intervalClient.organization) {
           intervalClient.#log.error('No organization defined')
           return
@@ -1212,6 +1260,10 @@ export default class IntervalClient {
         this.#closeTransaction(transactionId)
       },
       OPEN_PAGE: async inputs => {
+        if (this.#resolveShutdown) {
+          return { type: 'ERROR' as const, message: 'Host shutting down.' }
+        }
+
         if (!this.organization) {
           this.#log.error('No organization defined')
 
@@ -1627,6 +1679,12 @@ export default class IntervalClient {
             }
           }, this.#completeHttpRequestDelayMs)
         }
+
+        if (this.#resolveShutdown && this.#ioResponseHandlers.size === 0) {
+          setTimeout(() => {
+            this.#resolveShutdown?.()
+          }, this.#completeShutdownDelayMs)
+        }
       },
     }
   }

src/examples/shutdown/index.ts

@@ -0,0 +1,33 @@
+import Interval from '../../index'
+import { sleep } from '../utils/helpers'
+
+const interval = new Interval({
+  apiKey: 'live_N47qd1BrOMApNPmVd0BiDZQRLkocfdJKzvt8W6JT5ICemrAN',
+  logLevel: 'debug',
+  endpoint: 'ws://localhost:3000/websocket',
+  routes: {
+    wait_a_while: async (io, ctx) => {
+      await ctx.loading.start('Waiting...')
+      await sleep(5000)
+      return 'Done!'
+    },
+  },
+})
+
+interval.listen()
+
+process.on('SIGINT', () => {
+  interval
+    .safelyClose()
+    .then(() => {
+      console.log('Shut down!')
+      process.exit(0)
+    })
+    .catch(err => {
+      console.error(
+        'Failed shutting down gracefully, forcibly closing connection'
+      )
+      interval.immediatelyClose()
+      process.exit(0)
+    })
+})

src/experimental.ts

@@ -132,7 +132,7 @@ class ExperimentalInterval extends Interval {
     const client = new IntervalClient(this, this.config)
     const response = await client.respondToRequest(requestId)
 
-    client.close()
+    client.immediatelyClose()
 
     return response
   }
@@ -141,7 +141,7 @@ class ExperimentalInterval extends Interval {
     const client = new IntervalClient(this, this.config)
     const response = await client.declareHost(httpHostId)
 
-    client.close()
+    client.immediatelyClose()
 
     return response
   }

src/index.ts

@@ -231,8 +231,12 @@ export default class Interval {
     return this.#client.listen()
   }
 
-  close() {
-    return this.#client?.close()
+  immediatelyClose() {
+    return this.#client?.immediatelyClose()
+  }
+
+  async safelyClose(): Promise<void> {
+    return this.#client?.safelyClose()
   }
 
   /* @internal */ get client() {
@@ -278,7 +282,7 @@ export default class Interval {
       return
     }
 
-    let body: z.infer<typeof NOTIFY['inputs']>
+    let body: z.infer<(typeof NOTIFY)['inputs']>
     try {
       body = NOTIFY.inputs.parse({
         ...config,

src/internalRpcSchema.ts

@@ -414,6 +414,20 @@ export const wsServerSchema = {
       ])
       .nullable(),
   },
+  BEGIN_HOST_SHUTDOWN: {
+    // intentional empty object to allow for future additions
+    // and to support current `skipClientCall` behavior
+    inputs: z.object({}),
+    returns: z.discriminatedUnion('type', [
+      z.object({
+        type: z.literal('success'),
+      }),
+      z.object({
+        type: z.literal('error'),
+        message: z.string().optional(),
+      }),
+    ]),
+  },
 }
 
 export type WSServerSchema = typeof wsServerSchema
@@ -585,5 +599,5 @@ export const hostSchema = {
 export type HostSchema = typeof hostSchema
 
 export type PeerConnectionInitializer = (
-  inputs: z.infer<typeof INITIALIZE_PEER_CONNECTION['inputs']>
-) => Promise<z.infer<typeof INITIALIZE_PEER_CONNECTION['returns']>>
+  inputs: z.infer<(typeof INITIALIZE_PEER_CONNECTION)['inputs']>
+) => Promise<z.infer<(typeof INITIALIZE_PEER_CONNECTION)['returns']>>