React Native

/*
 * A WebSocket client for the RUAL cluster, for React Native.
 *
 * React Native ships a global WebSocket, so this file needs no dependency
 * beyond AppState, which is part of react-native itself.
 *
 * The socket carries a request/response protocol on top of plain WebSocket
 * frames. You send {method, uri, requestid} and the reply comes back carrying
 * the same requestid, so several requests can be in flight at once. The server
 * also pushes broadcasts that were never requested.
 */

import { AppState } from 'react-native';

const DEFAULT_OPTIONS = {
  // The socket lives in rual-core itself, on the cluster domain.
  //
  // Note the path is /ws. An older generation of this example used a separate
  // rual-ws host with a /primus path; that service is gone. If you still have
  // that URL somewhere, it will fail in a way that is easy to misread: the
  // hostname usually has no record of its own, resolves through the wildcard
  // to a load balancer, and gets a 301 redirect. A WebSocket handshake cannot
  // follow a redirect, so the socket never opens and nothing reports an error.
  'url': 'wss://{{cluster-domain}}/ws',

  // Reconnect backoff. Doubles from min to max.
  'reconnectMinMs': 500,
  'reconnectMaxMs': 30000,

  // How long a connection must stay up before its backoff is forgiven. See
  // scheduleReconnect for why this is not simply reset when the socket opens.
  'stableAfterMs': 10000,

  // How long to wait for a reply before rejecting the request.
  'requestTimeoutMs': 15000
};

/**
 * A WebSocket client for the RUAL cluster.
 *
 * @param {object} [options]
 * @param {string} [options.accessToken]      Token sent with every request.
 * @param {string} [options.url]              Defaults to the cluster socket.
 * @param {number} [options.reconnectMinMs]   First reconnect delay.
 * @param {number} [options.reconnectMaxMs]   Ceiling for the backoff.
 * @param {number} [options.stableAfterMs]    Survival time that forgives it.
 * @param {number} [options.requestTimeoutMs] Reply deadline for a request.
 */
export default class WS {
  constructor (options = {}) {
    this.options = Object.assign({}, DEFAULT_OPTIONS, options);
    this.accessToken = options.accessToken || null;

    this._connection = null;
    this.socketId = null;
    this.requestId = 0;
    this.openRequests = new Map();

    this.reconnecting = false;
    this.reconnectAttempts = 0;
    this.reconnectTimer = null;
    this.stableTimer = null;
    this.closedByUs = false;

    // Broadcast handlers, keyed by the broadcast payload type.
    this.listeners = new Map();

    // React Native only. See handleAppStateChange.
    this.appState = AppState.currentState;
    this.appStateSubscription = null;
  }

  /* ---------------------------------------------------------------- connect */

  /**
   * Open the connection and keep it open.
   *
   * Safe to call more than once: a socket that is already open, or still
   * connecting, is left alone. There is no need to await it, because
   * request() waits for the connection by itself.
   *
   * @returns {void}
   */
  connect () {
    if (this._connection) {
      const state = this._connection.readyState;
      if ((state === WebSocket.CONNECTING) || (state === WebSocket.OPEN)) {
        return;
      }
    }

    this.watchAppState();

    this.closedByUs = false;
    this.requestId = 0;
    this._connection = new WebSocket(this.options.url);

    this._connection.onopen = () => {
      console.log('[ws] connected');

      /*
       * Forgive the backoff only once this connection has STAYED up.
       *
       * A backend that is draining will accept you, assign an id, and then
       * immediately announce that it is going down. If the backoff resets
       * whenever a socket opens, it never grows, and you reconnect in a tight
       * loop against a server that is trying to shut down.
       */
      clearTimeout(this.stableTimer);
      this.stableTimer = setTimeout(() => {
        this.reconnectAttempts = 0;
      }, this.options.stableAfterMs);

      this.onOpen();
    };

    this._connection.onmessage = (event) => {
      this.handleMessage(event);
    };

    this._connection.onerror = () => {
      // Do not merely drop the socket. Without a scheduled reconnect nothing
      // reopens it, and because broadcasts are pushed rather than polled, the
      // client goes quiet instead of going wrong. That is much harder to spot.
      this.scheduleReconnect();
    };

    this._connection.onclose = () => {
      this.onClose();
      this.scheduleReconnect();
    };
  }

  /**
   * Tear the client down for good and stop it reconnecting. Call this when you
   * are finished with it, for example from the cleanup function of the effect
   * that owns it.
   *
   * Note that the socket handlers no longer call this. An error or a close now
   * schedules a reconnect instead, because a client that quietly stops
   * reconnecting is the defect this file exists to avoid.
   *
   * Requests still in flight reject with WS_DISCONNECTED.
   *
   * @returns {void}
   */
  disconnect () {
    this.closedByUs = true;

    if (this.appStateSubscription) {
      this.appStateSubscription.remove();
      this.appStateSubscription = null;
    }

    clearTimeout(this.reconnectTimer);
    this.reconnectTimer = null;
    this.reconnecting = false;

    this.teardown();
  }

  teardown () {
    clearTimeout(this.stableTimer);
    this.stableTimer = null;

    if (this._connection) {
      /*
       * Detach the handlers BEFORE closing.
       *
       * Otherwise close() fires our own onclose, which schedules a reconnect,
       * and a socket we deliberately tore down comes back from the dead. Worse,
       * the old instance keeps its own reconnect loop running alongside the new
       * one and the duplicates multiply.
       */
      const socket = this._connection;
      socket.onopen = null;
      socket.onmessage = null;
      socket.onerror = null;
      socket.onclose = null;
      try {
        socket.close();
      } catch (err) {
        // Already closing or closed. Nothing to do.
      }
    }

    this._connection = null;
    this.socketId = null;

    // Nothing can answer these now. Rejecting them immediately beats making
    // every caller sit out its own timeout, and it empties openRequests so the
    // request counter can safely start from zero on the next connection.
    this.failOpenRequests(new Error('WS_DISCONNECTED'));
  }

  failOpenRequests (err) {
    this.openRequests.forEach((pending) => {
      clearTimeout(pending.timer);
      pending.reject(err);
    });
    this.openRequests.clear();
  }

  /*
   * One reconnect cycle at a time, backed off exponentially and capped.
   *
   * The guard matters: close and error can both fire for the same failure, and
   * without it each would schedule its own timer, each timer would open its own
   * socket, and you would end up with several live connections all reconnecting.
   */
  scheduleReconnect () {
    if (this.closedByUs || this.reconnecting) {
      return;
    }
    this.reconnecting = true;

    this.teardown();

    this.reconnectAttempts += 1;
    const delay = Math.min(
      this.options.reconnectMinMs * Math.pow(2, (this.reconnectAttempts - 1)),
      this.options.reconnectMaxMs
    );
    console.log('[ws] reconnecting in ' + delay + 'ms (attempt ' + this.reconnectAttempts + ')');

    this.reconnectTimer = setTimeout(() => {
      this.reconnectTimer = null;
      this.reconnecting = false;
      if (!this.closedByUs) {
        this.connect();
      }
    }, delay);
  }

  async waitForConnected () {
    if (this._connection && (this._connection.readyState === WebSocket.OPEN)) {
      return;
    }

    if (!this._connection && !this.reconnecting) {
      this.connect();
    }

    // Poll rather than recurse: an unbounded recursive wait keeps a stack
    // frame per attempt and never gives up.
    const deadline = Date.now() + this.options.requestTimeoutMs;
    while (Date.now() < deadline) {
      if (this._connection && (this._connection.readyState === WebSocket.OPEN)) {
        return;
      }
      await new Promise((resolve) => {
        setTimeout(resolve, 100);
      });
    }

    throw new Error('WS_NOT_CONNECTED');
  }

  /* ------------------------------------------------------------- app state */

  /*
   * React Native only, and the reason this file is not just a copy of the
   * plain JavaScript client.
   *
   * When the app goes to the background the OS suspends the JavaScript thread
   * and generally kills the socket. Two things follow. The onclose handler may
   * not run while the thread is suspended, and a reconnect that was already
   * scheduled does not fire until the app is in the foreground again. A client
   * that reacts only to socket events therefore comes back to the foreground
   * holding a dead socket while still believing it is connected.
   *
   * Checking on the transition back to active is what makes the app feel
   * immediate rather than stale for the first few seconds.
   *
   * AppState.addEventListener returns a removable subscription on React Native
   * 0.65 and later. On older versions, use AppState.removeEventListener.
   */
  watchAppState () {
    if (this.appStateSubscription) {
      return;
    }

    this.appStateSubscription = AppState.addEventListener('change', (nextState) => {
      this.handleAppStateChange(nextState);
    });
  }

  handleAppStateChange (nextState) {
    const previous = this.appState;
    this.appState = nextState;

    if (this.closedByUs) {
      return;
    }
    if ((nextState !== 'active') || (previous === 'active')) {
      return;
    }
    if (this._connection && (this._connection.readyState === WebSocket.OPEN)) {
      return;
    }

    /*
     * The user just opened the app, so reconnect now instead of waiting out a
     * backoff that was measured against a backend problem. The socket was
     * almost certainly closed by the operating system rather than by a failing
     * server, which is why the attempt counter is cleared as well.
     */
    console.log('[ws] app returned to the foreground without a live socket');
    clearTimeout(this.reconnectTimer);
    this.reconnectTimer = null;
    this.reconnecting = false;
    this.reconnectAttempts = 0;

    this.teardown();
    this.connect();
  }

  /* ---------------------------------------------------------------- request */

  /**
   * Send a request and resolve with its result.
   *
   *   const me = await ws.request('GET', 'v1/users/me');
   *   await ws.request('POST', 'v1/some/endpoint', { body: { name: 'x' } });
   *
   * Every failure rejects with an Error, so one try/catch covers them all.
   * The Error carries the detail the server sent:
   *
   *   err.message  the server error, for example INVALID_ACCESS_TOKEN
   *   err.code     the status the server replied with
   *   err.uri      the uri that failed
   *   err.method   the method that failed
   *
   * Three failures come from the client rather than the server:
   *
   *   WS_NOT_CONNECTED    the socket did not come up in requestTimeoutMs
   *   WS_REQUEST_TIMEOUT  the reply did not arrive in requestTimeoutMs
   *   WS_DISCONNECTED     the socket dropped while this was in flight
   *
   * @param   {string} method                  GET, POST, PUT, DELETE.
   * @param   {string} uri                     Path after /api/, for example
   *                                           'v1/users/me'.
   * @param   {object} [options]
   * @param   {object} [options.body]          JSON body for the request.
   * @param   {object} [options.query]         Query parameters.
   * @param   {string} [options.accessToken]   Overrides the client token.
   * @returns {Promise<*>} The result the server returned.
   */
  async request (method, uri, options = {}) {
    await this.waitForConnected();

    this.requestId += 1;
    const requestid = 'RN' + String(this.requestId).padStart(10, '0');

    const payload = {
      'method': String(method).toUpperCase(),
      'requestid': requestid,
      'uri': '/api/' + uri
    };

    const token = options.accessToken || this.accessToken;
    if (token) {
      payload.access_token = token;
    }
    if (options.body) {
      payload.body = options.body;
    }
    if (options.query) {
      payload.query = options.query;
    }

    this._connection.send(JSON.stringify(payload));

    return new Promise((resolve, reject) => {
      // Always time out. Without this, a reply that never arrives leaks a
      // pending promise and its entry in openRequests, forever.
      const timer = setTimeout(() => {
        this.openRequests.delete(requestid);
        reject(new Error('WS_REQUEST_TIMEOUT'));
      }, this.options.requestTimeoutMs);

      this.openRequests.set(requestid, {
        'resolve': resolve,
        'reject': reject,
        'method': method,
        'uri': uri,
        'timer': timer
      });
    });
  }

  /**
   * The positional form this example used previously. Kept so existing code
   * carries on working, and because it reads well for a one off call.
   *
   * @param   {string} method
   * @param   {string} uri
   * @param   {object} [body]
   * @param   {object} [query]
   * @param   {string} [token]
   * @returns {Promise<*>} As request().
   */
  requestURI (method, uri, body = null, query = null, token = null) {
    return this.request(method, uri, {
      'body': body,
      'query': query,
      'accessToken': token
    });
  }

  /* -------------------------------------------------------------- broadcast */

  /**
   * Register a handler for a broadcast type.
   *
   * Handlers live on the client rather than on the socket, so they survive a
   * reconnect and do not have to be registered again.
   *
   * @param   {string}   type    The broadcast payload type to listen for.
   * @param   {function} handler Called with the broadcast payload.
   * @returns {function} Removes the handler again.
   */
  on (type, handler) {
    if (!this.listeners.has(type)) {
      this.listeners.set(type, new Set());
    }
    this.listeners.get(type).add(handler);

    return () => {
      const handlers = this.listeners.get(type);
      if (handlers) {
        handlers.delete(handler);
      }
    };
  }

  emit (type, payload) {
    const handlers = this.listeners.get(type);
    if (!handlers) {
      return;
    }
    handlers.forEach((handler) => {
      try {
        handler(payload);
      } catch (err) {
        console.error('[ws] broadcast handler for ' + type + ' threw', err);
      }
    });
  }

  /* ---------------------------------------------------------------- inbound */

  handleMessage (message) {
    if (!message.data) {
      return;
    }

    let data = null;
    try {
      data = JSON.parse(message.data);
    } catch (err) {
      console.error('[ws] received a frame that is not JSON');
      return;
    }

    // Heartbeat. The server sends a bare string and expects it echoed back
    // with "ping" swapped for "pong".
    if (typeof data === 'string') {
      this.replyPing(message.data);
      return;
    }

    // The server greets every new socket with the id it assigned. Keep it:
    // when broadcasts go missing, being able to say which socket you are on is
    // the difference between a five minute diagnosis and an hour of guessing.
    if (data.type === '_id') {
      this.socketId = data.id;
      console.log('[ws] socket id ' + data.id);
      return;
    }

    // The backend announces its own shutdown before it drops you. Acting on it
    // means reconnecting to whichever node comes up next, rather than waiting
    // for a TCP close that can take considerably longer to notice.
    if (data.server === 'down') {
      console.log('[ws] server announced it is going down');
      this.scheduleReconnect();
      return;
    }

    if (data.requestid) {
      this.resolveRequest(data);
      return;
    }

    if (data.broadcast) {
      const items = Array.isArray(data.broadcast) ? data.broadcast : [data.broadcast];
      items.forEach((item) => {
        const payload = item.message || item || {};
        this.emit(payload.type, payload);
      });
    }
  }

  resolveRequest (data) {
    const pending = this.openRequests.get(data.requestid);
    if (!pending) {
      // Usually a reply to a request that already timed out.
      console.log('[ws] reply for unknown request ' + data.requestid);
      return;
    }

    clearTimeout(pending.timer);
    this.openRequests.delete(data.requestid);

    if ((data.code === 200) || (data.code === 201)) {
      pending.resolve(data.result);
      return;
    }

    // Reject with a real Error so a stack trace survives, and hang the server
    // detail off it rather than throwing a bare object.
    const result = data.result || {};
    const err = new Error(result.error || ('HTTP_' + data.code));
    err.code = data.code;
    err.result = result;
    err.uri = pending.uri;
    err.method = pending.method;
    pending.reject(err);
  }

  replyPing (raw) {
    if (!this._connection || (this._connection.readyState !== WebSocket.OPEN)) {
      return;
    }
    this._connection.send(raw.replace('ping', 'pong'));
  }

  /* ------------------------------------------------------------ overridable */

  /**
   * Called once the socket is open. Assign your own implementation and
   * subscribe to what you need here: it runs again after every reconnect,
   * which is exactly when subscriptions have to be re-established.
   */
  onOpen () {}

  /**
   * Called whenever the socket closes, for any reason. Assign your own
   * implementation to drive connection state in the interface.
   */
  onClose () {
    console.log('[ws] disconnected');
  }
}

/*
 * Usage
 *
 *   const ws = new WS({ accessToken: myToken });
 *
 *   ws.onOpen = async () => {
 *     // Runs again after every reconnect, which is exactly when a
 *     // subscription has to be re-established.
 *     await ws.request('GET', 'v1/ws/subscribe');
 *   };
 *
 *   ws.on('some-broadcast-type', (payload) => {
 *     console.log('received', payload);
 *   });
 *
 *   ws.connect();
 *
 *   const me = await ws.request('GET', 'v1/users/me');
 *
 * Own it from an effect so that it is torn down when the screen goes away,
 * and so that Fast Refresh cannot leave a second client reconnecting behind
 * the first one:
 *
 *   useEffect(() => {
 *     const ws = new WS({ accessToken: myToken });
 *     ws.connect();
 *     return () => {
 *       ws.disconnect();
 *     };
 *   }, []);
 */

/*
 * Wiring the socket into a React Native app.
 *
 * One client is owned by the whole app rather than one per screen. A socket is
 * a connection, not a query: opening a second one costs another handshake,
 * another subscribe round, and doubles the broadcasts you receive.
 *
 * Put this provider near the root of the tree, above anything that needs data:
 *
 *   <SocketProvider accessToken={token}>
 *     <AppNavigator />
 *   </SocketProvider>
 */

import { createContext, useCallback, useContext, useEffect, useMemo, useRef, useState } from 'react';

import WS from './websocket';

const SocketContext = createContext(null);

export function SocketProvider ({ accessToken, children }) {
  /*
   * Create the client during the first render, not inside an effect.
   *
   * Child effects run before parent effects, so a screen that registers a
   * broadcast handler on mount would otherwise find no client and have its
   * handler quietly dropped. Building it here means the client exists before
   * any child can ask for it. The lazy form of useState runs the factory once,
   * however many times this component renders.
   */
  const [client] = useState(() => {
    return new WS();
  });

  const [connected, setConnected] = useState(false);

  /*
   * Keep the token current without dropping the socket.
   *
   * Tokens are refreshed while the app is running. Rebuilding the client on
   * every refresh would cost a reconnect and a fresh round of subscriptions
   * for no reason, so update the token that later requests will carry and
   * leave the connection alone.
   */
  useEffect(() => {
    client.accessToken = accessToken;
  }, [client, accessToken]);

  useEffect(() => {
    client.onOpen = async () => {
      setConnected(true);

      /*
       * Subscriptions belong to a socket, not to your session. After a
       * reconnect the server has no record of what the previous socket was
       * listening to, and onOpen runs again on every reconnect, which is
       * exactly what makes this the right place for it.
       */
      try {
        await client.request('GET', 'v1/ws/subscribe');
      } catch (err) {
        console.log('[socket] subscribe failed:', err.message);
      }
    };

    client.onClose = () => {
      setConnected(false);
    };

    client.connect();

    /*
     * Tearing down on unmount is what keeps Fast Refresh and StrictMode
     * honest. Without it, each remount would leave another client running its
     * own reconnect loop behind the new one.
     */
    return () => {
      client.disconnect();
    };
  }, [client]);

  /*
   * Send a request over the socket. See the screen example for the shape of
   * the errors it rejects with.
   */
  const request = useCallback((method, uri, options = {}) => {
    return client.request(method, uri, options);
  }, [client]);

  /*
   * Register a broadcast handler. Returns the function that removes it again,
   * which is what lets useBroadcast hand it straight back to React.
   *
   * Handlers live on the client rather than on the socket, so they survive a
   * reconnect and do not need re-registering.
   */
  const subscribe = useCallback((type, handler) => {
    return client.on(type, handler);
  }, [client]);

  const value = useMemo(() => {
    return {
      'connected': connected,
      'request': request,
      'subscribe': subscribe
    };
  }, [connected, request, subscribe]);

  return (
    <SocketContext.Provider value={value}>
      {children}
    </SocketContext.Provider>
  );
}

/*
 * Read the socket from any component below the provider.
 *
 *   const { connected, request } = useSocket();
 */
export function useSocket () {
  const context = useContext(SocketContext);
  if (!context) {
    throw new Error('useSocket must be used inside a SocketProvider');
  }

  return context;
}

/*
 * Listen to one broadcast type for as long as the component is mounted.
 *
 *   useBroadcast('device-updated', (payload) => {
 *     console.log(payload.guid, 'changed');
 *   });
 *
 * The handler is held in a ref and read at call time, so passing a new inline
 * function on every render does not tear the subscription down and set it up
 * again. Only the type has to stay stable.
 */
export function useBroadcast (type, handler) {
  const { subscribe } = useSocket();
  const saved = useRef(handler);

  useEffect(() => {
    saved.current = handler;
  }, [handler]);

  useEffect(() => {
    const unsubscribe = subscribe(type, (payload) => {
      saved.current(payload);
    });

    return unsubscribe;
  }, [subscribe, type]);
}

import { SocketProvider } from './socket-provider';

export default function App ({ accessToken }) {
  return (
    <SocketProvider accessToken={accessToken}>
      <AppNavigator />
    </SocketProvider>
  );
}

/*
 * Using the socket in a screen.
 *
 * This screen shows the three things almost every connected screen needs: an
 * initial load, live updates pushed from the server, and something sensible on
 * the display while the connection is down.
 */

import { useCallback, useEffect, useState } from 'react';
import { ActivityIndicator, FlatList, StyleSheet, Text, View } from 'react-native';

import { useBroadcast, useSocket } from './socket-provider';

export default function DevicesScreen () {
  const { connected, request } = useSocket();

  const [devices, setDevices] = useState([]);
  const [error, setError] = useState(null);
  const [loading, setLoading] = useState(true);

  const load = useCallback(async () => {
    setError(null);

    try {
      const result = await request('GET', 'v1/devices');
      setDevices(result.items || []);
    } catch (err) {
      /*
       * err.message is the error the server returned, for example
       * INVALID_ACCESS_TOKEN. err.code carries the status, and err.uri and
       * err.method describe the request that failed.
       *
       * The client rejects with WS_REQUEST_TIMEOUT when a reply never
       * arrives, and with WS_DISCONNECTED when the socket drops while the
       * request is still in flight. Neither leaves a promise pending.
       */
      setError(err.message);
    } finally {
      setLoading(false);
    }
  }, [request]);

  /*
   * Load when the socket is up, and load again after every reconnect.
   *
   * This is the part that is easy to leave out. Broadcasts that were sent
   * while the socket was down were missed, and nothing replays them, so the
   * copy held in state may be stale in ways no incoming message will correct.
   * Fetching again is the only way to be sure.
   */
  useEffect(() => {
    if (!connected) {
      return;
    }

    load();
  }, [connected, load]);

  const applyUpdate = useCallback((payload) => {
    setDevices((current) => {
      return current.map((device) => {
        if (device.guid !== payload.guid) {
          return device;
        }

        return { ...device, ...payload };
      });
    });
  }, []);

  useBroadcast('device-updated', applyUpdate);

  if (loading) {
    return (
      <View style={styles.centered}>
        <ActivityIndicator />
      </View>
    );
  }

  return (
    <View style={styles.container}>
      {/*
        * Say what is happening rather than showing stale data as though it
        * were live. The client is already reconnecting on its own, so this is
        * a status line and not a retry button.
        */}
      {connected === false && (
        <View style={styles.banner}>
          <Text style={styles.bannerText}>Reconnecting</Text>
        </View>
      )}

      {error !== null && (
        <Text style={styles.error}>{error}</Text>
      )}

      <FlatList
        data={devices}
        keyExtractor={(device) => {
          return device.guid;
        }}
        renderItem={({ item }) => {
          return (
            <View style={styles.row}>
              <Text style={styles.name}>{item.name}</Text>
              <Text style={styles.state}>{item.state}</Text>
            </View>
          );
        }}
      />
    </View>
  );
}

const styles = StyleSheet.create({
  'container': { 'flex': 1 },
  'centered': { 'flex': 1, 'alignItems': 'center', 'justifyContent': 'center' },
  'banner': { 'backgroundColor': '#8a6d3b', 'padding': 8 },
  'bannerText': { 'color': '#ffffff', 'textAlign': 'center' },
  'error': { 'color': '#a94442', 'padding': 12 },
  'row': { 'flexDirection': 'row', 'justifyContent': 'space-between', 'padding': 12 },
  'name': { 'fontWeight': '600' },
  'state': { 'opacity': 0.7 }
});