OkHttp SSE Coroutines 模块解读-携程实现

实现协程版的 EventSourceAsync

本文将详细解析如何实现 OkHttp SSE 的协程扩展 EventSourceAsync.kt，将传统的回调式 API 转换为现代的协程式 API。

实现目标

我们的目标是创建一个协程友好的 API，使开发者能够以更简洁、更现代的方式处理 Server-Sent Events (SSE)。具体来说，我们要实现三个主要功能：

将所有 SSE 事件转换为 Flow
提供只接收消息的 Channel
提供只接收消息数据的 Flow

前置知识

在开始实现之前，需要了解以下概念：

Server-Sent Events (SSE) - 一种基于 HTTP 的单向通信技术
OkHttp 的 EventSource API - 基于回调的 SSE 客户端实现
Kotlin 协程 - 特别是 Flow 和 Channel API
回调转换模式 - 将回调式 API 转换为协程式 API 的常用模式

实现步骤

步骤 1: 定义事件模型

首先，我们需要定义一个表示 SSE 事件的数据模型。我们使用密封类（sealed class）来表示不同类型的事件：

kotlin 复制代码

sealed class Event {
  data class Open(val response: Response) : Event()
  data class Message(val id: String?, val type: String?, val data: String) : Event()
  object Closed : Event()
  data class Error(val throwable: Throwable, val response: Response? = null) : Event()
}

这个模型包含四种事件类型：

Open: 连接建立时触发
Message: 接收到服务器消息时触发
Closed: 连接正常关闭时触发
Error: 发生错误时触发

步骤 2: 实现 events() 函数

events() 函数是最基础的 API，它返回一个包含所有类型事件的 Flow：

kotlin 复制代码

fun EventSources.Factory.events(request: Request): Flow<Event> = callbackFlow {
  val eventSource = newEventSource(
    request = request,
    listener = object : EventSourceListener() {
      override fun onOpen(eventSource: EventSource, response: Response) {
        trySend(Event.Open(response))
      }

      override fun onEvent(
        eventSource: EventSource,
        id: String?,
        type: String?,
        data: String,
      ) {
        trySend(Event.Message(id, type, data))
      }

      override fun onClosed(eventSource: EventSource) {
        trySend(Event.Closed)
        close()
      }

      override fun onFailure(
        eventSource: EventSource,
        throwable: Throwable?,
        response: Response?,
      ) {
        if (throwable != null) {
          trySend(Event.Error(throwable, response))
        }
        close()
      }
    },
  )

  awaitClose {
    eventSource.cancel()
  }
}

实现解析：

使用 callbackFlow ：callbackFlow 是 Kotlin 协程库提供的构建器，专门用于将回调式 API 转换为 Flow。它创建一个 Flow，同时提供一个 ProducerScope 用于发送元素。
创建 EventSource ：使用 OkHttp 的 newEventSource 方法创建一个 EventSource，并传入自定义的 EventSourceListener。
实现回调方法：
- onOpen: 当连接建立时，发送 Event.Open 事件
- onEvent: 当接收到消息时，发送 Event.Message 事件
- onClosed: 当连接关闭时，发送 Event.Closed 事件并关闭 Flow
- onFailure: 当发生错误时，发送 Event.Error 事件并关闭 Flow
使用 awaitClose ：awaitClose 是 callbackFlow 的一个函数，它会在 Flow 被取消或关闭时执行提供的代码块。在这里，我们使用它来取消 EventSource，确保资源被正确释放。

步骤 3: 实现 messages() 函数

messages() 函数返回一个只包含消息事件的 Channel：

kotlin 复制代码

suspend fun EventSources.Factory.messages(request: Request): ReceiveChannel<Event.Message> {
  val channel = Channel<Event.Message>(Channel.UNLIMITED)

  val eventSource = newEventSource(
    request = request,
    listener = object : EventSourceListener() {
      override fun onEvent(
        eventSource: EventSource,
        id: String?,
        type: String?,
        data: String,
      ) {
        channel.trySend(Event.Message(id, type, data))
      }

      override fun onClosed(eventSource: EventSource) {
        channel.close()
      }

      override fun onFailure(
        eventSource: EventSource,
        throwable: Throwable?,
        response: Response?,
      ) {
        if (throwable != null) {
          channel.close(throwable)
        } else {
          channel.close()
        }
      }
    },
  )

  return channel
}

实现解析：

创建 Channel ：使用 Channel.UNLIMITED 创建一个无限容量的 Channel，避免背压问题。
创建 EventSource ：与 events() 类似，但这里我们只关注消息事件。
实现回调方法：
- onEvent: 将消息发送到 Channel
- onClosed: 关闭 Channel
- onFailure: 如果有错误，使用错误关闭 Channel；否则正常关闭
返回 Channel：直接返回创建的 Channel，让调用者可以接收消息。

注意，这个函数没有使用 callbackFlow，而是直接使用了 Channel API。这是因为 Channel 提供了更直接的控制，特别是在处理错误时。

步骤 4: 实现 messageData() 函数

messageData() 函数返回一个只包含消息数据字符串的 Flow：

kotlin 复制代码

fun EventSources.Factory.messageData(request: Request): Flow<String> = callbackFlow {
  val eventSource = newEventSource(
    request = request,
    listener = object : EventSourceListener() {
      override fun onEvent(
        eventSource: EventSource,
        id: String?,
        type: String?,
        data: String,
      ) {
        trySend(data)
      }

      override fun onClosed(eventSource: EventSource) {
        close()
      }

      override fun onFailure(
        eventSource: EventSource,
        throwable: Throwable?,
        response: Response?,
      ) {
        if (throwable != null) {
          close(throwable)
        } else {
          close()
        }
      }
    },
  )

  awaitClose {
    eventSource.cancel()
  }
}

实现解析：

使用 callbackFlow ：与 events() 类似，但这里我们只发送消息数据字符串。
实现回调方法：
- onEvent: 只发送 data 字符串，忽略 ID 和类型
- onClosed: 关闭 Flow
- onFailure: 如果有错误，使用错误关闭 Flow；否则正常关闭
使用 awaitClose：同样确保在 Flow 取消时释放资源。

关键技术点解析

1. callbackFlow 的使用

callbackFlow 是将回调转换为 Flow 的关键工具。它创建一个 Flow，同时提供一个作用域，在这个作用域中可以：

使用 trySend() 发送元素
使用 close() 关闭 Flow
使用 awaitClose {} 注册清理代码

kotlin 复制代码

fun someCallbackApi(): Flow<Result> = callbackFlow {
  val callback = object : Callback {
    override fun onSuccess(result: Result) {
      trySend(result)
    }

    override fun onComplete() {
      close()
    }

    override fun onError(error: Exception) {
      close(error)
    }
  }

  // 注册回调
  api.registerCallback(callback)

  // 当 Flow 被取消时执行
  awaitClose {
    api.unregisterCallback(callback)
  }
}

2. Channel 与 Flow 的选择

在 messages() 函数中，我们选择返回 Channel 而不是 Flow。这是因为：

Channel 提供了更直接的控制，特别是在处理错误时
Channel 可以在多个协程之间共享
如果需要，可以使用 receiveAsFlow() 将 Channel 转换为 Flow

如果你更喜欢统一的 API，可以将 messages() 也实现为返回 Flow：

kotlin 复制代码

fun EventSources.Factory.messages(request: Request): Flow<Event.Message> = events(request)
  .filterIsInstance<Event.Message>()

但这种实现会丢失一些错误处理的灵活性。

3. 错误处理策略

在协程 API 中，错误处理有两种主要策略：

传播错误 ：使用 close(throwable) 或 throw exception，让调用者处理错误
转换为事件 ：将错误包装为事件（如 Event.Error），让调用者决定如何处理

在我们的实现中：

events() 使用了第二种策略，将错误转换为 Event.Error 事件
messages() 和 messageData() 使用了第一种策略，直接传播错误

选择哪种策略取决于 API 的设计目标和使用场景。

4. 资源管理

在协程 API 中，资源管理是一个重要问题。我们使用 awaitClose 确保在 Flow 被取消时释放资源：

kotlin 复制代码

awaitClose {
  eventSource.cancel()
}

这确保了即使在异常情况下，资源也能被正确释放，避免泄漏。

完整实现

将上述所有部分组合起来，我们就得到了完整的 EventSourceAsync.kt 实现：

kotlin 复制代码

sealed class Event {
  data class Open(val response: Response) : Event()
  data class Message(val id: String?, val type: String?, val data: String) : Event()
  object Closed : Event()
  data class Error(val throwable: Throwable, val response: Response? = null) : Event()
}

fun EventSources.Factory.events(request: Request): Flow<Event> = callbackFlow {
  val eventSource = newEventSource(
    request = request,
    listener = object : EventSourceListener() {
      override fun onOpen(eventSource: EventSource, response: Response) {
        trySend(Event.Open(response))
      }

      override fun onEvent(
        eventSource: EventSource,
        id: String?,
        type: String?,
        data: String,
      ) {
        trySend(Event.Message(id, type, data))
      }

      override fun onClosed(eventSource: EventSource) {
        trySend(Event.Closed)
        close()
      }

      override fun onFailure(
        eventSource: EventSource,
        throwable: Throwable?,
        response: Response?,
      ) {
        if (throwable != null) {
          trySend(Event.Error(throwable, response))
        }
        close()
      }
    },
  )

  awaitClose {
    eventSource.cancel()
  }
}

suspend fun EventSources.Factory.messages(request: Request): ReceiveChannel<Event.Message> {
  val channel = Channel<Event.Message>(Channel.UNLIMITED)

  val eventSource = newEventSource(
    request = request,
    listener = object : EventSourceListener() {
      override fun onEvent(
        eventSource: EventSource,
        id: String?,
        type: String?,
        data: String,
      ) {
        channel.trySend(Event.Message(id, type, data))
      }

      override fun onClosed(eventSource: EventSource) {
        channel.close()
      }

      override fun onFailure(
        eventSource: EventSource,
        throwable: Throwable?,
        response: Response?,
      ) {
        if (throwable != null) {
          channel.close(throwable)
        } else {
          channel.close()
        }
      }
    },
  )

  return channel
}

fun EventSources.Factory.messageData(request: Request): Flow<String> = callbackFlow {
  val eventSource = newEventSource(
    request = request,
    listener = object : EventSourceListener() {
      override fun onEvent(
        eventSource: EventSource,
        id: String?,
        type: String?,
        data: String,
      ) {
        trySend(data)
      }

      override fun onClosed(eventSource: EventSource) {
        close()
      }

      override fun onFailure(
        eventSource: EventSource,
        throwable: Throwable?,
        response: Response?,
      ) {
        if (throwable != null) {
          close(throwable)
        } else {
          close()
        }
      }
    },
  )

  awaitClose {
    eventSource.cancel()
  }
}

扩展与改进

这个基本实现可以根据需要进行扩展和改进：

1. 添加重试机制

kotlin 复制代码

fun EventSources.Factory.eventsWithRetry(
  request: Request,
  maxRetries: Int = 3,
  delayMillis: Long = 1000
): Flow<Event> = flow {
  var retries = 0
  while (true) {
    try {
      events(request).collect { event ->
        emit(event)
        if (event is Event.Closed) {
          return@flow
        }
      }
    } catch (e: IOException) {
      if (retries >= maxRetries) {
        throw e
      }
      retries++
      delay(delayMillis * retries)
    }
  }
}

2. 添加超时机制

kotlin 复制代码

fun EventSources.Factory.eventsWithTimeout(
  request: Request,
  timeoutMillis: Long = 30000
): Flow<Event> = events(request)
  .timeout(timeoutMillis) {
    throw TimeoutException("EventSource connection timed out after $timeoutMillis ms")
  }

3. 添加背压处理

kotlin 复制代码

fun EventSources.Factory.eventsWithBackpressure(
  request: Request,
  capacity: Int = 64,
  onBufferOverflow: BufferOverflow = BufferOverflow.SUSPEND
): Flow<Event> = events(request)
  .buffer(capacity, onBufferOverflow)

4. 添加事件过滤和转换

kotlin 复制代码

// 只获取特定类型的消息
fun EventSources.Factory.messagesOfType(
  request: Request,
  type: String
): Flow<Event.Message> = events(request)
  .filterIsInstance<Event.Message>()
  .filter { it.type == type }

// 将消息转换为特定类型的对象
inline fun <reified T> EventSources.Factory.messagesAs(
  request: Request,
  crossinline transform: (Event.Message) -> T
): Flow<T> = events(request)
  .filterIsInstance<Event.Message>()
  .map { transform(it) }

实现中的注意事项

1. 线程安全

Kotlin 协程库确保了 Flow 和 Channel 的线程安全，但在与回调 API 交互时，需要注意线程问题。在我们的实现中，trySend() 和 close() 方法是线程安全的，可以从任何线程调用。

2. 内存泄漏防护

确保在不再需要 EventSource 时取消它是很重要的。我们的实现通过 awaitClose 确保了这一点，但调用者也应该确保在适当的生命周期内使用这些 API。

kotlin 复制代码

// 在 ViewModel 中使用
val job = viewModelScope.launch {
  factory.events(request).collect { /* 处理事件 */ }
}

// 在 onCleared 中取消
override fun onCleared() {
  job.cancel()
  super.onCleared()
}

3. 错误处理最佳实践

在使用这些 API 时，应该始终处理错误：

kotlin 复制代码

// 使用 events() API
factory.events(request)
  .catch { error ->
    // 处理错误
    Log.e("SSE", "Error in SSE connection", error)
    emit(Event.Error(error))
  }
  .collect { event ->
    // 处理事件
  }

// 使用 messageData() API
factory.messageData(request)
  .catch { error ->
    // 处理错误
    Log.e("SSE", "Error in SSE connection", error)
  }
  .collect { data ->
    // 处理数据
  }

总结

实现协程版的 EventSourceAsync 主要涉及以下步骤：

定义表示 SSE 事件的数据模型
使用 callbackFlow 将回调式 API 转换为 Flow
实现不同级别的抽象，满足不同的使用需求
确保正确的资源管理和错误处理

通过这种方式，我们成功地将 OkHttp 的回调式 SSE API 转换为更现代、更简洁的协程式 API，使得处理服务器发送的事件变得更加直观和高效。

这种模式不仅适用于 SSE，也可以应用于其他回调式 API，如 WebSocket、蓝牙通信等。掌握这种转换模式，可以帮助你将任何回调式 API 转换为协程式 API，提高代码的可读性和可维护性。