golem_rust/websocket.rs
1// Copyright 2024-2026 Golem Cloud
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::bindings::golem::websocket::client::{
16 Error, Message, WebsocketConnection as RawWebsocketConnection,
17};
18
19pub use crate::bindings::golem::websocket::client::{
20 CloseInfo as WebSocketCloseInfo, Error as WebSocketError, Message as WebSocketMessage,
21 ReconstructionPolicy as WebSocketReconstructionPolicy,
22};
23
24/// A WebSocket connection with both blocking and async receive methods.
25pub struct WebsocketConnection {
26 inner: RawWebsocketConnection,
27}
28
29impl WebsocketConnection {
30 /// Connect to a WebSocket server at the given URL (ws:// or wss://).
31 /// Optional headers for auth, subprotocols, etc.
32 pub fn connect(
33 url: &str,
34 headers: Option<Vec<(String, String)>>,
35 reconstruction_policy: Option<WebSocketReconstructionPolicy>,
36 ) -> Result<Self, Error> {
37 RawWebsocketConnection::connect(url, headers.as_deref(), reconstruction_policy)
38 .map(|inner| Self { inner })
39 }
40
41 /// Send a message (text or binary).
42 pub fn send(&self, message: &Message) -> Result<(), Error> {
43 self.inner.send(message)
44 }
45
46 /// Receive the next message, blocking until one is available.
47 pub fn blocking_receive(&self) -> Result<Message, Error> {
48 wit_bindgen::block_on(async move { self.inner.receive().await })
49 }
50
51 /// Receive the next message, blocking with a timeout in milliseconds.
52 /// Returns `None` if the timeout expires before a message arrives.
53 pub fn blocking_receive_with_timeout(&self, timeout_ms: u64) -> Result<Option<Message>, Error> {
54 wit_bindgen::block_on(async move { self.inner.receive_with_timeout(timeout_ms).await })
55 }
56
57 /// Receive the next message asynchronously.
58 /// Yields the current task until a message is available.
59 pub async fn receive(&self) -> Result<Message, Error> {
60 self.inner.receive().await
61 }
62
63 /// Receive the next message asynchronously with a timeout in milliseconds.
64 /// Yields the current task; returns `None` if the timeout expires before a message arrives.
65 pub async fn receive_with_timeout(&self, timeout_ms: u64) -> Result<Option<Message>, Error> {
66 self.inner.receive_with_timeout(timeout_ms).await
67 }
68
69 /// Send a close frame with optional code and reason.
70 pub fn close(&self, code: Option<u16>, reason: Option<String>) -> Result<(), Error> {
71 self.inner.close(code, reason.as_deref())
72 }
73}