pleroma-fe/src/stores/streaming.js

import { defineStore } from 'pinia'

import { useOAuthStore } from 'src/stores/oauth.js'

import {
  getMastodonSocketURI,
  ProcessedWS,
  WSConnectionStatus,
} from 'src/api/websocket.js'

const ARGUMENT_MAP = {
  tag: 'tag',
  list: 'list',
}

export const TIMELINE_STREAM_MAP = {
  friends: 'user',
  public: 'public',
  tag: 'hashtag',
  list: 'list',
  dms: 'direct',
}

const retryTimeout = (multiplier) => 1000 * multiplier

export class StreamStateEvent extends Event {
  original

  constructor(name, original) {
    super(name)
    this.original = original
  }
}

export class StreamErrorEvent extends Event {
  error

  constructor(error) {
    super('error', error)
    this.error = error
  }
}

export class StreamMessageEvent extends Event {
  data
  stream
  timestamp

  constructor(name, stream, data) {
    super(name)
    this.data = data
    this.stream = stream
    this.timestamp = Date.now()
  }
}

export const useStreamingStore = defineStore('streaming', {
  state: () => ({
    socket: null,
    error: null,
    state: null,
    retryMultiplier: 1,
    retrying: false,
    subscribers: new Set(),
    subscriptions: new Map(),
    globalSubscriptions: new Set(),
  }),
  actions: {
    addSubscriber(subscriber) {
      const { stream, et } = subscriber

      if (stream) {
        if (!this.subscriptions.has(stream.name)) {
          this.subscriptions.set(stream.name, new Map())
        }

        const streamSubs = this.subscriptions.get(stream.name)

        if (streamSubs.has(stream.argument)) {
          throw new Error('Subscription already exists!')
        }

        streamSubs.set(stream.argument, subscriber)
      } else {
        this.globalSubscriptions.add(subscriber)
      }

      this.subscribers.add(subscriber)
      if (this.state === WSConnectionStatus.JOINED) {
        if (stream) {
          this.socket.subscribe(...this.getSubArgs(stream))
        }
        et.dispatchEvent(new StreamStateEvent('open'))
      }
    },
    removeSubscriber(subscriber) {
      const { stream } = subscriber

      this.subscribers.delete(subscriber)
      if (stream) {
        this.subscriptions.get(stream.name).delete(stream.argument)
      }

      if (this.state === WSConnectionStatus.JOINED) {
        this.socket.unsubscribe(...this.getSubArgs(stream))
      }
    },
    initSocket(initial) {
      this.state = initial
        ? WSConnectionStatus.STARTING_INITIAL
        : WSConnectionStatus.STARTING

      const credentials = useOAuthStore().token
      const url = getMastodonSocketURI({ credentials })

      this.socket = ProcessedWS({
        url,
        id: 'Unified',
        credentials,
      })

      this.socket.addEventListener('pleroma:authenticated', this.onAuth)
      this.socket.addEventListener('open', this.onOpen)
      this.socket.addEventListener('close', this.onClose)
      this.socket.addEventListener('message', this.onMessage)
      this.socket.addEventListener('error', this.onError)
    },
    stopSocket() {
      this.socket.close()
      this.state = WSConnectionStatus.CLOSED
    },

    getSubArgs(stream) {
      const argumentKey = ARGUMENT_MAP[stream.name]
      const args = argumentKey
        ? {
            [argumentKey]: stream.argument,
          }
        : null

      return [stream.name, args]
    },
    onAuth() {
      this.subscribers.forEach(({ stream, et }) => {
        et.dispatchEvent(new StreamStateEvent('authenticated'))

        if (stream) {
          this.socket.subscribe(...this.getSubArgs(stream))
        }
      })
      this.state = WSConnectionStatus.JOINED
    },
    onOpen() {
      this.retryMultiplier = 1
      this.retrying = false
      this.error = null
      this.subscribers.forEach(({ stream, et }) => {
        et.dispatchEvent(new StreamStateEvent('open'))
      })
    },
    onMessage({ data: message }) {
      if (!message) return // pings
      const { event: eventName, stream: eventStream, ...data } = message
      const [streamName, streamArgument] = eventStream ?? []

      const subscriber = this.subscriptions.get(streamName)?.get(streamArgument)

      const totalSubs = [
        ...this.globalSubscriptions.values(),
        subscriber,
      ].filter(Boolean)

      const eventData = (() => {
        switch (eventName) {
          case 'status.update':
          case 'update':
            return [data.status]
          case 'notification':
            return [data.notification]
          case 'delete':
            return [data.id]
          default:
            return data
        }
      })()

      const event = new StreamMessageEvent(
        eventName,
        { name: streamName, argument: streamArgument },
        eventData,
      )

      totalSubs.forEach(({ stream, et }) => {
        et.dispatchEvent(event)
      })
    },
    onError({ data: error }) {
      this.subscribers.forEach(({ stream, et }) => {
        et.dispatchEvent(new StreamErrorEvent(error))
      })
      this.error = error
      console.error('Error in MastoAPI websocket:', error)
    },
    onClose({ data: closeEvent }) {
      const ignoreCodes = new Set([
        1000, // Normal (intended) closure
        1001, // Going away
      ])
      const { code } = closeEvent

      if (ignoreCodes.has(code)) {
        console.debug(
          `Not restarting socket becasue of closure code ${code} is in ignore list`,
        )

        this.state = WSConnectionStatus.CLOSED
        this.retrying = false
        this.error = null
        this.retryMultiplier = 1

        this.subscribers.forEach(({ et }) => {
          et.dispatchEvent(new StreamStateEvent('close', closeEvent))
        })
      } else {
        console.warn(
          `MastoAPI websocket disconnected, restarting. CloseEvent code: ${code}`,
        )

        setTimeout(() => {
          this.initSocket()
        }, retryTimeout(this.retryMultiplier))

        this.retryMultiplier += 1

        if (!this.retrying) {
          this.subscribers.forEach(({ et }) => {
            et.dispatchEvent(new StreamStateEvent('close', closeEvent))
          })
        }

        this.retrying = true
        this.state = WSConnectionStatus.ERROR
      }
    },
  },
})