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}