Server-Sent Events

import Vapor

/// One Server-Sent Event, serialised exactly as the WHATWG "event stream" format describes.
struct ServerSentEvent: Sendable {
    var id: String?
    var event: String?
    var data: String
    var retry: Int?          // milliseconds the browser should wait before reconnecting
    var comment: String?     // a line starting with ":" is ignored by EventSource: a keep-alive

    /// A comment-only frame used as a heartbeat.
    static func keepAlive(_ text: String = "keep-alive") -> ServerSentEvent {
        ServerSentEvent(data: "", comment: text)
    }

    /// The bytes written to the socket. A blank line ends the event.
    var wireFormat: String {
        var out = ""
        if let comment {
            for line in Self.lines(comment) { out += ": \(line)\n" }
        }
        if let id { out += "id: \(Self.singleLine(id))\n" }
        if let event { out += "event: \(Self.singleLine(event))\n" }
        if let retry { out += "retry: \(retry)\n" }
        // each line of the payload needs its own "data:" prefix; a frame without data dispatches nothing
        if !data.isEmpty {
            for line in Self.lines(data) { out += "data: \(line)\n" }
        }
        return out + "\n"
    }

    /// Splits on LF, CR and CRLF, the three line endings the format allows. Swift treats "\r\n" as ONE Character,
    /// so the split works on unicode scalars.
    private static func lines(_ text: String) -> [String] {
        var result: [String] = []
        var current = String.UnicodeScalarView()
        var previousWasCR = false
        for scalar in text.unicodeScalars {
            switch scalar {
            case "\r":
                result.append(String(current)); current = String.UnicodeScalarView(); previousWasCR = true
            case "\n":
                if previousWasCR { previousWasCR = false } else { result.append(String(current)); current = String.UnicodeScalarView() }
            default:
                previousWasCR = false; current.append(scalar)
            }
        }
        result.append(String(current))
        return result
    }

    /// `id` and `event` are single-line fields: drop any line break rather than let it inject a field.
    private static func singleLine(_ value: String) -> String {
        String(String.UnicodeScalarView(value.unicodeScalars.filter { $0 != "\n" && $0 != "\r" }))
    }
}

extension Response {
    /// Builds a `text/event-stream` response around a producer closure.
    static func eventStream(
        _ produce: @escaping @Sendable (_ send: @Sendable (ServerSentEvent) async throws -> Void) async throws -> Void
    ) -> Response {
        var headers = HTTPHeaders()
        headers.replaceOrAdd(name: .contentType, value: "text/event-stream; charset=utf-8")
        headers.replaceOrAdd(name: .cacheControl, value: "no-cache, no-transform")
        headers.replaceOrAdd(name: "X-Accel-Buffering", value: "no")   // nginx: do not buffer this response
        let body = Response.Body(managedAsyncStream: { writer in
            let send: @Sendable (ServerSentEvent) async throws -> Void = { event in
                try await writer.writeBuffer(ByteBuffer(string: event.wireFormat))
            }
            try await produce(send)
        })
        return Response(status: .ok, headers: headers, body: body)
    }
}

import Vapor

struct LoanDueNotice: Content, Sendable {
    let loanID: UUID
    let bookTitle: String
    let dueAt: Date
}

/// What a subscriber receives: a numbered notice, or a tick from the heartbeat task.
enum HubItem: Sendable {
    case notice(id: Int, LoanDueNotice)
    case heartbeat
}

/// Fan-out of notices to every open stream, with a bounded replay buffer for Last-Event-ID.
actor NotificationHub {
    private var nextID = 1
    private var buffer: [(id: Int, notice: LoanDueNotice)] = []
    private let capacity: Int
    private var subscribers: [UUID: AsyncStream<HubItem>.Continuation] = [:]

    init(capacity: Int = 100) { self.capacity = capacity }

    var subscriberCount: Int { subscribers.count }

    func publish(_ notice: LoanDueNotice) {
        let id = nextID
        nextID += 1
        buffer.append((id, notice))
        if buffer.count > capacity { buffer.removeFirst(buffer.count - capacity) }
        for continuation in subscribers.values { continuation.yield(.notice(id: id, notice)) }
    }

    /// Returns the missed notices (ids greater than `lastID`) and a live stream. Both are produced inside one actor
    /// turn, so no notice can fall between the replay and the live part.
    func subscribe(after lastID: Int?) -> (replay: [HubItem], live: AsyncStream<HubItem>, token: UUID) {
        let token = UUID()
        let replay = buffer.filter { $0.id > (lastID ?? Int.max) }.map { HubItem.notice(id: $0.id, $0.notice) }
        let (live, continuation) = AsyncStream.makeStream(of: HubItem.self, bufferingPolicy: .bufferingNewest(256))
        subscribers[token] = continuation
        return (replay, live, token)
    }

    func heartbeat(_ token: UUID) {
        subscribers[token]?.yield(.heartbeat)
    }

    /// Ends every live stream so the producers return and the connections close.
    func finishAll() {
        for continuation in subscribers.values { continuation.finish() }
        subscribers.removeAll()
    }

    func unsubscribe(_ token: UUID) {
        subscribers.removeValue(forKey: token)?.finish()
    }
}

struct HubKey: StorageKey { typealias Value = NotificationHub }

extension Application {
    var hub: NotificationHub {
        get { storage[HubKey.self]! }
        set { storage[HubKey.self] = newValue }
    }
}

import Fluent
import Vapor

func routes(_ app: Application) throws {
    let notifications = app.grouped("notifications")

    // GET /notifications/stream
    notifications.get("stream") { req async throws -> Response in
        let lastID = req.headers.first(name: "Last-Event-ID").flatMap(Int.init)
        let hub = req.application.hub
        let logger = req.logger
        let interval = Environment.get("SSE_HEARTBEAT").flatMap(Int.init) ?? 15
        let (replay, live, token) = await hub.subscribe(after: lastID)

        return Response.eventStream { send in
            defer { Task { await hub.unsubscribe(token); logger.info("subscriber removed") } }
            let heartbeat = Task {
                while !Task.isCancelled {
                    try? await Task.sleep(for: .seconds(interval))
                    await hub.heartbeat(token)
                }
            }
            defer { heartbeat.cancel() }

            try await send(ServerSentEvent(data: "", retry: 3000, comment: "stream open"))
            for item in replay { try await send(item.frame) }
            for await item in live {
                try await send(item.frame)
            }
            logger.info("stream finished")
        }
    }

    notifications.get("subscribers") { req async -> [String: Int] in
        ["open": await req.application.hub.subscriberCount]
    }

    // POST /notifications/scan : publish a notice for every loan due within a day
    notifications.post("scan") { req async throws -> [String: Int] in
        let horizon = Date().addingTimeInterval(24 * 3600)
        let loans = try await Loan.query(on: req.db)
            .filter(\.$returnedAt == nil)
            .filter(\.$dueAt <= horizon)
            .with(\.$book)
            .all()
        for loan in loans {
            await req.application.hub.publish(
                LoanDueNotice(loanID: try loan.requireID(), bookTitle: loan.book.title, dueAt: loan.dueAt))
        }
        return ["published": loans.count]
    }
}

extension HubItem {
    var frame: ServerSentEvent {
        switch self {
        case .notice(let id, let notice):
            let json = (try? String(data: JSONEncoder.custom(dates: .iso8601).encode(notice), encoding: .utf8)) ?? "{}"
            return ServerSentEvent(id: String(id), event: "loan-due", data: json)
        case .heartbeat:
            return .keepAlive()
        }
    }
}

[ ERROR ] Could not stop HTTP server: Abort.500: Server stop took too long.

/// Without this, an open stream keeps the server from stopping until the shutdown timeout expires.
struct FinishStreamsOnShutdown: LifecycleHandler {
    let hub: NotificationHub
    func shutdownAsync(_ application: Application) async {
        await hub.finishAll()
    }
}

import Fluent
import FluentSQLiteDriver
import Vapor

func configure(_ app: Application) async throws {
    app.databases.use(.sqlite(.memory), as: .sqlite)
    app.migrations.add(CreateBookshelf())
    try await app.autoMigrate()
    app.hub = NotificationHub()
    // registered after the HTTP server's own handler, so it runs first on shutdown
    app.lifecycle.use(FinishStreamsOnShutdown(hub: app.hub))
    try routes(app)
}

const source = new EventSource("/notifications/stream");

source.addEventListener("loan-due", (e) => {
  const notice = JSON.parse(e.data);
  showToast(`"${notice.bookTitle}" is due ${new Date(notice.dueAt).toLocaleString()}`);
});

source.onerror = () => {
  // EventSource reconnects by itself while readyState is CONNECTING;
  // CLOSED means the server answered with a non-200 or the wrong content type
  if (source.readyState === EventSource.CLOSED) console.warn("stream closed for good");
};