Skip to main content

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}