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

实现协程版的 EventSourceAsync

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

实现目标

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

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

前置知识

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

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

实现步骤

步骤 1: 定义事件模型

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

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：

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()
  }
}

实现解析：

1. 使用 callbackFlow：callbackFlow 是 Kotlin 协程库提供的构建器，专门用于将回调式 API 转换为 Flow。它创建一个 Flow，同时提供一个 ProducerScope 用于发送元素。

2. 创建 EventSource：使用 OkHttp 的 newEventSource 方法创建一个 EventSource，并传入自定义的 EventSourceListener。

3. 实现回调方法：

   • onOpen: 当连接建立时，发送 Event.Open 事件
   • onEvent: 当接收到消息时，发送 Event.Message 事件
   • onClosed: 当连接关闭时，发送 Event.Closed 事件并关闭 Flow
   • onFailure: 当发生错误时，发送 Event.Error 事件并关闭 Flow

4. 使用 awaitClose：awaitClose 是 callbackFlow 的一个函数，它会在 Flow 被取消或关闭时执行提供的代码块。在这里，我们使用它来取消 EventSource，确保资源被正确释放。

步骤 3: 实现 messages() 函数

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

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
}

实现解析：

1. 创建 Channel：使用 Channel.UNLIMITED 创建一个无限容量的 Channel，避免背压问题。

2. 创建 EventSource：与 events() 类似，但这里我们只关注消息事件。

3. 实现回调方法：

   • onEvent: 将消息发送到 Channel
   • onClosed: 关闭 Channel
   • onFailure: 如果有错误，使用错误关闭 Channel；否则正常关闭

4. 返回 Channel：直接返回创建的 Channel，让调用者可以接收消息。

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

步骤 4: 实现 messageData() 函数

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

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. 使用 callbackFlow：与 events() 类似，但这里我们只发送消息数据字符串。

2. 实现回调方法：

   • onEvent: 只发送 data 字符串，忽略 ID 和类型
   • onClosed: 关闭 Flow
   • onFailure: 如果有错误，使用错误关闭 Flow；否则正常关闭

3. 使用 awaitClose：同样确保在 Flow 取消时释放资源。

关键技术点解析

1. callbackFlow 的使用

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

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

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。这是因为：

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

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

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

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

3. 错误处理策略

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

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

在我们的实现中：

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

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

4. 资源管理

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

awaitClose {
  eventSource.cancel()
}

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

完整实现

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

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. 添加重试机制

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. 添加超时机制

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. 添加背压处理

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

4. 添加事件过滤和转换

// 只获取特定类型的消息
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。

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

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

3. 错误处理最佳实践

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

// 使用 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 主要涉及以下步骤：

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

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

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