mass-subscriptions.js

#!/usr/bin/env node

const path = require('path')
const { parse } = require('graphql')
const { ApolloClient } = require('apollo-client')
const { WebSocketLink } = require('apollo-link-ws')
const { SubscriptionClient } = require('subscriptions-transport-ws')
const { InMemoryCache } = require('apollo-cache-inmemory')
const ws = require('ws')

let args = process.argv

if (args.length < 5) {
  console.error(`Usage: ${path.basename(args[1])} <node-url> <number-of-connections> <number-of-subscriptions> '<query>'`)
  process.exit(1)
}

const NODE_URL = args[2]
const CONNECTIONS = parseInt(args[3])
const SUBSCRIPTIONS = parseInt(args[4])
const QUERY = args[5]

const clients = []

for (let conn = 0; conn < CONNECTIONS; conn++) {
  let label = `Connection ${conn}/${CONNECTIONS-1}`
  
  console.log(`${label}: Connecting`)

  let wsClient = new SubscriptionClient(NODE_URL, {
    reconnect: false,
    connectionParams: {
    }
  }, ws)
  
  wsClient.onConnected(() => console.log(`${label}: Connected`))
  wsClient.onError((e) => console.error(`${label}: Failed to connect. Error: ${e.message}`))
  // wsClient.onDisconnected(() => console.log(`${label}: Disconnected`))
    
  let client = new ApolloClient({ 
    link: new WebSocketLink(wsClient),
    cache: new InMemoryCache(),
  })
  
  
  clients[conn] = client
  
  let subscriptions = new Set()
  let updates = []
  
  for (let i = 0; i < Math.ceil(SUBSCRIPTIONS / CONNECTIONS); i++) {
    let subscription = client.subscribe({
      query: parse(QUERY),
    })
    
    subscription.subscribe({
      next(data) {
        updates.push(i)
      },
      error(e) {
        console.error(`${label}: Failed: ${e.message}`)
      },
    })
    
    subscriptions.add(subscription)
  }
  
  setInterval(() => {
    if (updates.length > 0) {
      console.log(
        `${label} (${Math.ceil(SUBSCRIPTIONS / CONNECTIONS)} subscriptions): ` +
          `${updates.length} subscription updates`
      )
    }
    updates = []
  }, 2000)
}