feat(data): Events Lib WS Publishing (#3027)

Author: tylerjroach

appsync/aws-appsync-events/src/main/java/com/amplifyframework/aws/appsync/events/Events.kt

@@ -122,7 +122,7 @@ class Events @VisibleForTesting internal constructor(
     fun channel(
         channelName: String,
         authorizers: ChannelAuthorizers = this.defaultChannelAuthorizers,
-    ) = EventsChannel(channelName, authorizers, endpoints, eventsWebSocketProvider)
+    ) = EventsChannel(channelName, authorizers, eventsWebSocketProvider)
 
     /**
      * Method to disconnect from all channels.

appsync/aws-appsync-events/src/main/java/com/amplifyframework/aws/appsync/events/EventsChannel.kt

@@ -15,7 +15,6 @@
 package com.amplifyframework.aws.appsync.events
 
 import com.amplifyframework.aws.appsync.core.AppSyncAuthorizer
-import com.amplifyframework.aws.appsync.core.AppSyncRequest
 import com.amplifyframework.aws.appsync.events.data.ChannelAuthorizers
 import com.amplifyframework.aws.appsync.events.data.ConnectionClosedException
 import com.amplifyframework.aws.appsync.events.data.EventsException
@@ -34,7 +33,9 @@ import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.onCompletion
 import kotlinx.coroutines.flow.onStart
+import kotlinx.serialization.json.JsonArray
 import kotlinx.serialization.json.JsonElement
+import kotlinx.serialization.json.JsonPrimitive
 
 /**
  * A class to manage channel subscriptions and publishes
@@ -45,7 +46,6 @@ import kotlinx.serialization.json.JsonElement
 class EventsChannel internal constructor(
     val name: String,
     val authorizers: ChannelAuthorizers,
-    private val endpoints: EventsEndpoints,
     private val eventsWebSocketProvider: EventsWebSocketProvider
 ) {
 
@@ -83,12 +83,12 @@ class EventsChannel internal constructor(
      * @param authorizer for the publish call. If not provided, the EventChannel publish authorizer will be used.
      * @return result of publish.
      */
-    @Throws(EventsException::class)
+    @Throws(Exception::class)
     suspend fun publish(
         event: JsonElement,
         authorizer: AppSyncAuthorizer = this.authorizers.publishAuthorizer
     ): PublishResult {
-        TODO("Need to implement")
+        return publish(listOf(event), authorizer)
     }
 
     /**
@@ -98,12 +98,43 @@ class EventsChannel internal constructor(
      * @param authorizer for the publish call. If not provided, the EventChannel publish authorizer will be used.
      * @return result of publish.
      */
-    @Throws(EventsException::class)
+    @Throws(Exception::class)
     suspend fun publish(
         events: List<JsonElement>,
         authorizer: AppSyncAuthorizer = this.authorizers.publishAuthorizer
-    ): PublishResult {
-        TODO("Need to implement")
+    ): PublishResult = coroutineScope {
+        val publishId = UUID.randomUUID().toString()
+        val publishMessage = WebSocketMessage.Send.Publish(
+            id = publishId,
+            channel = name,
+            events = JsonArray(events.map { JsonPrimitive(it.toString()) }),
+        )
+
+        val webSocket = eventsWebSocketProvider.getConnectedWebSocket()
+        val deferredResponse = async { getPublishResponse(webSocket, publishId) }
+
+        val queued = webSocket.sendWithAuthorizer(publishMessage, authorizer)
+        if (!queued) {
+            throw webSocket.disconnectReason?.toCloseException() ?: ConnectionClosedException()
+        }
+
+        return@coroutineScope when (val response = deferredResponse.await()) {
+            is WebSocketMessage.Received.PublishSuccess -> {
+                PublishResult(response.successfulEvents, response.failedEvents)
+            }
+
+            is WebSocketMessage.ErrorContainer -> {
+                val fallbackMessage = "Failed to publish event(s)"
+                throw response.errors.firstOrNull()?.toEventsException(fallbackMessage)
+                    ?: EventsException(fallbackMessage)
+            }
+
+            is WebSocketMessage.Closed -> {
+                throw response.reason.toCloseException()
+            }
+
+            else -> throw EventsException("Received unexpected publish response of type: ${response::class}")
+        }
     }
 
     private fun createSubscriptionEventDataFlow(subscriptionHolder: SubscriptionHolder): Flow<EventsMessage> {
@@ -116,11 +147,12 @@ class EventsChannel internal constructor(
                         emit(EventsMessage(it.event))
                     }
                     it is WebSocketMessage.Closed -> {
-                        if (it.reason is DisconnectReason.UserInitiated) {
-                            throw UserClosedConnectionException()
-                        } else {
-                            throw ConnectionClosedException(it.reason.throwable)
-                        }
+                        throw it.reason.toCloseException()
+                    }
+                    it is WebSocketMessage.ErrorContainer && it.id == subscriptionHolder.id -> {
+                        val exceptionMessage = "Received error for subscription"
+                        throw it.errors.firstOrNull()?.toEventsException(exceptionMessage)
+                            ?: EventsException(exceptionMessage)
                     }
                     else -> Unit
                 }
@@ -133,24 +165,31 @@ class EventsChannel internal constructor(
         subscriptionId: String,
         authorizer: AppSyncAuthorizer
     ): Boolean = coroutineScope {
+        // create a deferred holder for subscription response
         val deferredSubscriptionResponse = async { getSubscriptionResponse(webSocket, subscriptionId) }
 
         // Publish subscription to websocket
-        publishSubscription(webSocket, subscriptionId, authorizer)
+        val queued = webSocket.sendWithAuthorizer(
+            webSocketMessage = WebSocketMessage.Send.Subscription.Subscribe(id = subscriptionId, channel = name),
+            authorizer = authorizer
+        )
+        if (!queued) {
+            throw webSocket.disconnectReason?.toCloseException() ?: ConnectionClosedException()
+        }
 
         // Wait for subscription result to return
         when (val response = deferredSubscriptionResponse.await()) {
             is WebSocketMessage.Received.Subscription.SubscribeSuccess -> {
                 return@coroutineScope true
             }
-            is WebSocketMessage.Received.Subscription.SubscribeError -> {
+            is WebSocketMessage.ErrorContainer -> {
                 val exceptionMessage = "Subscribe failed for channel: $name"
                 throw response.errors.firstOrNull()
                     ?.toEventsException(exceptionMessage)
                     ?: EventsException(exceptionMessage)
             }
-            is WebSocketMessage.Received.ConnectionClosed -> {
-                throw ConnectionClosedException()
+            is WebSocketMessage.Closed -> {
+                throw response.reason.toCloseException()
             }
             else -> throw EventsException("Received unexpected subscription response of type: ${response::class}")
         }
@@ -160,34 +199,22 @@ class EventsChannel internal constructor(
         return webSocket.events.first {
             when {
                 it is WebSocketMessage.Received.Subscription && it.id == subscriptionId -> true
-                it is WebSocketMessage.Received.ConnectionClosed -> true
+                it is WebSocketMessage.ErrorContainer && it.id == subscriptionId -> true
+                it is WebSocketMessage.Closed -> true
                 else -> false
             }
         }
     }
 
-    private suspend fun publishSubscription(
-        webSocket: EventsWebSocket,
-        subscriptionId: String,
-        authorizer: AppSyncAuthorizer
-    ) {
-        val subscribeMessage = WebSocketMessage.Send.Subscription.Subscribe(
-            id = subscriptionId,
-            channel = name,
-            authorization = authorizer.getAuthorizationHeaders(
-                object : AppSyncRequest {
-                    override val url: String
-                        get() = endpoints.restEndpoint.toString()
-                    override val body: String?
-                        get() = null
-                    override val headers: Map<String, String>
-                        get() = emptyMap()
-                    override val method: AppSyncRequest.HttpMethod
-                        get() = AppSyncRequest.HttpMethod.GET
-                }
-            )
-        )
-        webSocket.send(subscribeMessage)
+    private suspend fun getPublishResponse(webSocket: EventsWebSocket, publishId: String): WebSocketMessage {
+        return webSocket.events.first {
+            when {
+                it is WebSocketMessage.Received.PublishSuccess && it.id == publishId -> true
+                it is WebSocketMessage.ErrorContainer && it.id == publishId -> true
+                it is WebSocketMessage.Closed -> true
+                else -> false
+            }
+        }
     }
 
     private fun completeSubscription(subscriptionHolder: SubscriptionHolder, throwable: Throwable?) {
@@ -198,7 +225,11 @@ class EventsChannel internal constructor(
 
         if (currentWebSocket != null && isSubscribed && !isDisconnected) {
             // Unsubscribe from channel when flow is completed
-            currentWebSocket.send(WebSocketMessage.Send.Subscription.Unsubscribe(subscriptionHolder.id))
+            try {
+                currentWebSocket.send(WebSocketMessage.Send.Subscription.Unsubscribe(subscriptionHolder.id))
+            } catch (e: Exception) {
+                // do nothing with a failed unsubscribe post
+            }
         }
     }
 }

appsync/aws-appsync-events/src/main/java/com/amplifyframework/aws/appsync/events/EventsWebSocket.kt

@@ -31,6 +31,10 @@ import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.first
 import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.JsonObject
+import kotlinx.serialization.json.JsonPrimitive
+import kotlinx.serialization.json.encodeToJsonElement
+import kotlinx.serialization.json.jsonObject
 import okhttp3.OkHttpClient
 import okhttp3.Request
 import okhttp3.Response
@@ -50,8 +54,9 @@ internal class EventsWebSocket(
 
     private lateinit var webSocket: WebSocket
     @Volatile internal var isClosed = false
-    private var disconnectReason: DisconnectReason? = null
+    internal var disconnectReason: WebSocketDisconnectReason? = null
     private val connectionTimeoutTimer = ConnectionTimeoutTimer(onTimeout = ::onTimeout)
+    val preAuthPublishHeaders: Map<String, String> by lazy { mapOf(HeaderKeys.HOST to eventsEndpoints.host) }
     private val logger = loggerProvider?.getLogger(TAG)
 
     @Throws(ConnectException::class)
@@ -75,7 +80,7 @@ internal class EventsWebSocket(
                 webSocket.cancel()
                 throw ConnectException(connectionResponse.reason.throwable)
             }
-            is WebSocketMessage.Received.ConnectionError -> {
+            is WebSocketMessage.ErrorContainer -> {
                 webSocket.cancel()
                 throw ConnectException(
                     connectionResponse.errors.firstOrNull()?.toEventsException()
@@ -91,7 +96,7 @@ internal class EventsWebSocket(
     }
 
     suspend fun disconnect(flushEvents: Boolean) = coroutineScope {
-        disconnectReason = DisconnectReason.UserInitiated
+        disconnectReason = WebSocketDisconnectReason.UserInitiated
         val deferredClosedResponse = async { getClosedResponse() }
         when (flushEvents) {
             true -> webSocket.close(NORMAL_CLOSE_CODE, "User initiated disconnect")
@@ -101,16 +106,20 @@ internal class EventsWebSocket(
     }
 
     override fun onOpen(webSocket: WebSocket, response: Response) {
-        logger?.debug ("onOpen: sending connection init")
-        send(WebSocketMessage.Send.ConnectionInit())
+        logger?.debug("onOpen: sending connection init")
+        try {
+            send(WebSocketMessage.Send.ConnectionInit())
+        } catch (e: Exception) {
+            logger?.error("onOpen: exception encountered", e) // do nothing. closure will handle error
+        }
     }
 
     override fun onMessage(webSocket: WebSocket, text: String) {
         connectionTimeoutTimer.resetTimeoutTimer()
         logger?.debug { "Websocket onMessage: $text" }
         try {
             val eventMessage = json.decodeFromString<WebSocketMessage.Received>(text)
-            _events.tryEmit(eventMessage)
+            emitEvent(eventMessage)
         } catch (e: Exception) {
             logger?.error(e) { "Websocket onMessage: exception encountered" }
         }
@@ -128,27 +137,69 @@ internal class EventsWebSocket(
     override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
         // Events api sends normal close code even in failure
         // so inspecting code/reason isn't helpful as it should be
-        logger?.debug {"onClosed: reason = $disconnectReason" }
+        logger?.debug { "onClosed: reason = $disconnectReason" }
         handleClosed()
     }
 
     private fun onTimeout() {
-        disconnectReason = DisconnectReason.Timeout
+        disconnectReason = WebSocketDisconnectReason.Timeout
         webSocket.cancel()
     }
 
     private fun handleClosed() {
         connectionTimeoutTimer.stop()
-        _events.tryEmit(
-            WebSocketMessage.Closed(reason = disconnectReason ?: DisconnectReason.Service())
-        )
+        emitEvent(WebSocketMessage.Closed(reason = disconnectReason ?: WebSocketDisconnectReason.Service()))
         isClosed = true
     }
 
-    inline fun <reified T : WebSocketMessage> send(webSocketMessage: T) {
-        val message = json.encodeToString(webSocketMessage)
-        logger?.debug { "send: ${webSocketMessage::class.java}" }
-        webSocket.send(message)
+    // returns true if websocket queued up event. false if failed
+    internal suspend inline fun <reified T : WebSocketMessage.Send> sendWithAuthorizer(
+        webSocketMessage: T,
+        authorizer: AppSyncAuthorizer
+    ): Boolean {
+        // The base message will not include id, type, or authorization fields
+        val baseMessageJson = json.encodeToJsonElement(webSocketMessage).jsonObject
+
+        // Create the authorization headers first
+        val authHeaders = authorizer.getAuthorizationHeaders(object : AppSyncRequest {
+            override val method = AppSyncRequest.HttpMethod.POST
+            override val url = eventsEndpoints.restEndpoint.toString()
+            override val headers = preAuthPublishHeaders
+            override val body = json.encodeToString(baseMessageJson)
+        })
+
+        // We reconstruct the message, adding in the id, type, and authorization fields
+        val message = json.encodeToString(
+            JsonObject(
+                buildMap {
+                    putAll(baseMessageJson)
+                    put("id", JsonPrimitive(webSocketMessage.id))
+                    put("type", JsonPrimitive(webSocketMessage.type))
+                    put(
+                        "authorization",
+                        JsonObject((preAuthPublishHeaders + authHeaders).mapValues { JsonPrimitive(it.value) })
+                    )
+                }
+            )
+        )
+
+        return send(message)
+    }
+
+    // returns true if websocket queued up event. false if failed
+    inline fun <reified T : WebSocketMessage> send(webSocketMessage: T): Boolean {
+        return send(json.encodeToString(webSocketMessage))
+    }
+
+    // returns true if websocket queued up event. false if failed
+    private fun send(eventJson: String): Boolean {
+        logger?.debug { "send: $eventJson" }
+        return webSocket.send(eventJson)
+    }
+
+    private fun emitEvent(event: WebSocketMessage) {
+        logger?.debug { "emit ${event::class.java}" }
+        _events.tryEmit(event)
     }
 
     companion object {
@@ -206,9 +257,3 @@ private class ConnectAppSyncRequest(
     override val body: String
         get() = "{}"
 }
-
-internal sealed class DisconnectReason(val throwable: Throwable?) {
-    data object UserInitiated : DisconnectReason(null)
-    data object Timeout : DisconnectReason(EventsException("Connection timed out."))
-    class Service(throwable: Throwable? = null) : DisconnectReason(throwable)
-}

appsync/aws-appsync-events/src/main/java/com/amplifyframework/aws/appsync/events/EventsWebSocketProvider.kt

@@ -40,7 +40,6 @@ internal class EventsWebSocketProvider(
     val existingWebSocket: EventsWebSocket?
         get() = connectionResultReference.get()?.getOrNull()
 
-
     suspend fun getConnectedWebSocket(): EventsWebSocket = getConnectedWebSocketResult().getOrThrow()
 
     private suspend fun getConnectedWebSocketResult(): Result<EventsWebSocket> = coroutineScope {