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};
22
23/// A WebSocket connection with both blocking and async receive methods.
24pub struct WebsocketConnection {
25    inner: RawWebsocketConnection,
26}
27
28impl WebsocketConnection {
29    /// Connect to a WebSocket server at the given URL (ws:// or wss://).
30    /// Optional headers for auth, subprotocols, etc.
31    pub fn connect(url: &str, headers: Option<Vec<(String, String)>>) -> Result<Self, Error> {
32        RawWebsocketConnection::connect(url, headers.as_deref()).map(|inner| Self { inner })
33    }
34
35    /// Send a message (text or binary).
36    pub fn send(&self, message: &Message) -> Result<(), Error> {
37        self.inner.send(message)
38    }
39
40    /// Receive the next message, blocking until one is available.
41    pub fn blocking_receive(&self) -> Result<Message, Error> {
42        self.inner.receive()
43    }
44
45    /// Receive the next message, blocking with a timeout in milliseconds.
46    /// Returns `None` if the timeout expires before a message arrives.
47    pub fn blocking_receive_with_timeout(&self, timeout_ms: u64) -> Result<Option<Message>, Error> {
48        self.inner.receive_with_timeout(timeout_ms)
49    }
50
51    /// Receive the next message asynchronously.
52    /// Yields the current task until a message is available.
53    pub async fn receive(&self) -> Result<Message, Error> {
54        let pollable = self.inner.subscribe();
55        wstd::io::AsyncPollable::new(pollable).wait_for().await;
56        self.inner.receive()
57    }
58
59    /// Send a close frame with optional code and reason.
60    pub fn close(&self, code: Option<u16>, reason: Option<String>) -> Result<(), Error> {
61        self.inner.close(code, reason.as_deref())
62    }
63}