Export netimgui files in MuJoCo Copybara configuration for use with the upcoming Web Viewer
PiperOrigin-RevId: 954769526 Change-Id: I4cdee083082f7d55c2babd5e664ec0283a2aaf63
This commit is contained in:
Matija Kecman
committed by
Copybara-Service
Copybara-Service
parent
c7b6e0b8dd
commit
ca5337161d
@@ -0,0 +1,324 @@
|
||||
// Copyright 2026 DeepMind Technologies Limited
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// https://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
// WASM networking backend for NetImgui.
|
||||
//
|
||||
// Implements the NetImgui::Internal::Network interface using Emscripten
|
||||
// WebSockets. This file is compiled only under the Emscripten toolchain and
|
||||
// provides the browser-side transport layer for netimgui draw data.
|
||||
#include "NetImgui_Shared.h"
|
||||
|
||||
#ifndef __EMSCRIPTEN__
|
||||
#error "This file must be compiled with emscripten."
|
||||
#endif
|
||||
|
||||
#include <emscripten/console.h>
|
||||
#include <emscripten/websocket.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <atomic>
|
||||
#include <cstring>
|
||||
#include <mutex>
|
||||
#include <vector>
|
||||
|
||||
#include "NetImgui_CmdPackets.h"
|
||||
#include "google/logging.h"
|
||||
#include "google/network_status.h"
|
||||
|
||||
namespace NetImgui {
|
||||
namespace Internal {
|
||||
namespace Network {
|
||||
|
||||
struct SocketInfo {
|
||||
EMSCRIPTEN_WEBSOCKET_T mSocket = 0;
|
||||
std::atomic<bool> mConnected{false};
|
||||
std::atomic<bool> mError{false};
|
||||
std::atomic<bool> mClosed{false};
|
||||
std::atomic<int> mCloseCode{0};
|
||||
|
||||
std::vector<uint8_t> mBuffer;
|
||||
std::mutex mBufferMutex;
|
||||
int mSendSizeMax =
|
||||
1024 * 1024; // Interface compatibility with other backends
|
||||
};
|
||||
|
||||
// --- WebSocket event callbacks ---
|
||||
|
||||
static EM_BOOL OnWebSocketOpen(int /*event_type*/,
|
||||
const EmscriptenWebSocketOpenEvent* /*event*/,
|
||||
void* user_data) {
|
||||
auto* socket = static_cast<SocketInfo*>(user_data);
|
||||
if (socket) {
|
||||
socket->mConnected = true;
|
||||
}
|
||||
return EM_TRUE;
|
||||
}
|
||||
|
||||
static EM_BOOL OnWebSocketMessage(int /*event_type*/,
|
||||
const EmscriptenWebSocketMessageEvent* event,
|
||||
void* user_data) {
|
||||
auto* socket = static_cast<SocketInfo*>(user_data);
|
||||
if (socket && !event->isText) {
|
||||
std::lock_guard<std::mutex> lock(socket->mBufferMutex);
|
||||
size_t old_size = socket->mBuffer.size();
|
||||
socket->mBuffer.insert(socket->mBuffer.end(), event->data,
|
||||
event->data + event->numBytes);
|
||||
static int msg_count = 0;
|
||||
++msg_count;
|
||||
VLOG(1, "onmessage #%d: %d bytes, buffer: %zu -> %zu", msg_count,
|
||||
event->numBytes, old_size, socket->mBuffer.size());
|
||||
} else if (socket && event->isText) {
|
||||
VLOG(1, "onmessage: TEXT frame (%d bytes), IGNORED", event->numBytes);
|
||||
}
|
||||
return EM_TRUE;
|
||||
}
|
||||
|
||||
static EM_BOOL OnWebSocketClose(int /*event_type*/,
|
||||
const EmscriptenWebSocketCloseEvent* event,
|
||||
void* user_data) {
|
||||
auto* socket = static_cast<SocketInfo*>(user_data);
|
||||
if (socket) {
|
||||
socket->mClosed = true;
|
||||
socket->mCloseCode = event->code;
|
||||
}
|
||||
return EM_TRUE;
|
||||
}
|
||||
|
||||
static EM_BOOL OnWebSocketError(int /*event_type*/,
|
||||
const EmscriptenWebSocketErrorEvent* /*event*/,
|
||||
void* user_data) {
|
||||
auto* socket = static_cast<SocketInfo*>(user_data);
|
||||
if (socket) {
|
||||
socket->mError = true;
|
||||
}
|
||||
return EM_TRUE;
|
||||
}
|
||||
|
||||
// --- Network interface implementation ---
|
||||
|
||||
bool Startup() { return emscripten_websocket_is_supported(); }
|
||||
|
||||
void Shutdown() {}
|
||||
|
||||
SocketInfo* Connect(const char* server_host, uint32_t /*server_port*/) {
|
||||
if (!emscripten_websocket_is_supported()) return nullptr;
|
||||
|
||||
SocketInfo* socket_info = netImguiNew<SocketInfo>();
|
||||
|
||||
EmscriptenWebSocketCreateAttributes attr;
|
||||
emscripten_websocket_init_create_attributes(&attr);
|
||||
|
||||
// The web viewer always passes a complete WebSocket URL (e.g.
|
||||
// "ws://host:8080/ui", built from the page origin by WsUrl), so it is used
|
||||
// verbatim and server_port is ignored — a browser reaches the viewer's
|
||||
// paths on the page's own shared port, not a dedicated NetImgui port.
|
||||
const std::string url(server_host);
|
||||
attr.url = url.c_str();
|
||||
attr.createOnMainThread = EM_TRUE;
|
||||
LOG(Info, "Connecting WebSocket to: %s", url.c_str());
|
||||
|
||||
socket_info->mSocket = emscripten_websocket_new(&attr);
|
||||
if (socket_info->mSocket <= 0) {
|
||||
netImguiDelete(socket_info);
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
emscripten_websocket_set_onopen_callback(socket_info->mSocket, socket_info,
|
||||
OnWebSocketOpen);
|
||||
emscripten_websocket_set_onmessage_callback(socket_info->mSocket, socket_info,
|
||||
OnWebSocketMessage);
|
||||
emscripten_websocket_set_onclose_callback(socket_info->mSocket, socket_info,
|
||||
OnWebSocketClose);
|
||||
emscripten_websocket_set_onerror_callback(socket_info->mSocket, socket_info,
|
||||
OnWebSocketError);
|
||||
|
||||
return socket_info;
|
||||
}
|
||||
|
||||
// Abandoned sockets are not freed immediately: with -pthread, websocket
|
||||
// events are queued across threads, so an already-queued close/error event
|
||||
// can still dereference the SocketInfo after emscripten_websocket_delete().
|
||||
// When the freed block was recycled for the next socket, such a late event
|
||||
// stamped a stale mClosed flag onto a healthy connection, which the
|
||||
// reconnect logic then tore down — a self-sustaining reconnect loop. Keep
|
||||
// abandoned sockets in a small ring and free them several disconnects
|
||||
// later, when any queued events are long gone.
|
||||
static SocketInfo* s_socket_graveyard[8] = {};
|
||||
static int s_socket_graveyard_idx = 0;
|
||||
|
||||
void Disconnect(SocketInfo* client_socket) {
|
||||
if (!client_socket) return;
|
||||
client_socket->mClosed = true;
|
||||
// Detach this socket from future events; queued events may already hold
|
||||
// the pointer (the graveyard above covers those).
|
||||
emscripten_websocket_set_onopen_callback(client_socket->mSocket, nullptr,
|
||||
OnWebSocketOpen);
|
||||
emscripten_websocket_set_onmessage_callback(client_socket->mSocket, nullptr,
|
||||
OnWebSocketMessage);
|
||||
emscripten_websocket_set_onclose_callback(client_socket->mSocket, nullptr,
|
||||
OnWebSocketClose);
|
||||
emscripten_websocket_set_onerror_callback(client_socket->mSocket, nullptr,
|
||||
OnWebSocketError);
|
||||
emscripten_websocket_close(client_socket->mSocket, 1000,
|
||||
"Normal Disconnection");
|
||||
emscripten_websocket_delete(client_socket->mSocket);
|
||||
{
|
||||
// Release the receive buffer now; only the flags must stay valid.
|
||||
std::lock_guard<std::mutex> lock(client_socket->mBufferMutex);
|
||||
client_socket->mBuffer.clear();
|
||||
client_socket->mBuffer.shrink_to_fit();
|
||||
}
|
||||
if (s_socket_graveyard[s_socket_graveyard_idx]) {
|
||||
netImguiDelete(s_socket_graveyard[s_socket_graveyard_idx]);
|
||||
}
|
||||
s_socket_graveyard[s_socket_graveyard_idx] = client_socket;
|
||||
s_socket_graveyard_idx = (s_socket_graveyard_idx + 1) % 8;
|
||||
}
|
||||
|
||||
bool DataReceivePending(SocketInfo* client_socket) {
|
||||
if (!client_socket) return false;
|
||||
|
||||
if (client_socket->mError || client_socket->mClosed) {
|
||||
// Connection is dead — flush any buffered data so we stop processing
|
||||
// stale commands that arrived before the close.
|
||||
std::lock_guard<std::mutex> lock(client_socket->mBufferMutex);
|
||||
if (!client_socket->mBuffer.empty()) {
|
||||
LOG(Warning, "Connection closed/error. Discarding %zu buffered bytes.",
|
||||
client_socket->mBuffer.size());
|
||||
client_socket->mBuffer.clear();
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
std::lock_guard<std::mutex> lock(client_socket->mBufferMutex);
|
||||
return !client_socket->mBuffer.empty();
|
||||
}
|
||||
|
||||
void DataReceive(SocketInfo* client_socket, PendingCom& pending_rcv) {
|
||||
if (!client_socket || !pending_rcv.pCommand) {
|
||||
pending_rcv.bError = true;
|
||||
return;
|
||||
}
|
||||
|
||||
if (!client_socket->mConnected) {
|
||||
pending_rcv.bError = false; // Not ready yet, caller will retry.
|
||||
return;
|
||||
}
|
||||
|
||||
// The size field comes off the wire; a value smaller than what has already
|
||||
// been read (e.g. a command header claiming < 8 bytes, from a corrupted or
|
||||
// desynced stream) would underflow the subtraction below into a huge
|
||||
// size_t and memcpy past the destination command buffer.
|
||||
if (pending_rcv.pCommand->mSize < pending_rcv.SizeCurrent) {
|
||||
LOG(Error, "DataReceive: wire size %u < %zu bytes already read; stream "
|
||||
"is corrupt",
|
||||
pending_rcv.pCommand->mSize,
|
||||
static_cast<size_t>(pending_rcv.SizeCurrent));
|
||||
pending_rcv.bError = true;
|
||||
return;
|
||||
}
|
||||
|
||||
size_t bytes_to_read = pending_rcv.pCommand->mSize - pending_rcv.SizeCurrent;
|
||||
if (bytes_to_read == 0) return;
|
||||
|
||||
std::lock_guard<std::mutex> lock(client_socket->mBufferMutex);
|
||||
|
||||
VLOG(1, "DataReceive: want=%zu, have=%zu, cmd_size=%u, progress=%zu",
|
||||
bytes_to_read, client_socket->mBuffer.size(),
|
||||
pending_rcv.pCommand->mSize, pending_rcv.SizeCurrent);
|
||||
|
||||
if (client_socket->mBuffer.empty()) {
|
||||
if (client_socket->mError || client_socket->mClosed) {
|
||||
pending_rcv.bError = true;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
size_t bytes_to_consume =
|
||||
std::min(bytes_to_read, client_socket->mBuffer.size());
|
||||
if (bytes_to_consume > 0) {
|
||||
memcpy(reinterpret_cast<uint8_t*>(pending_rcv.pCommand) +
|
||||
pending_rcv.SizeCurrent,
|
||||
client_socket->mBuffer.data(), bytes_to_consume);
|
||||
client_socket->mBuffer.erase(
|
||||
client_socket->mBuffer.begin(),
|
||||
client_socket->mBuffer.begin() + bytes_to_consume);
|
||||
pending_rcv.SizeCurrent += bytes_to_consume;
|
||||
pending_rcv.bError = false;
|
||||
}
|
||||
}
|
||||
|
||||
void DataSend(SocketInfo* client_socket, PendingCom& pending_send) {
|
||||
if (!client_socket || client_socket->mClosed || client_socket->mError ||
|
||||
!pending_send.pCommand) {
|
||||
pending_send.bError = true;
|
||||
return;
|
||||
}
|
||||
|
||||
if (!client_socket->mConnected) {
|
||||
pending_send.bError = false; // Not ready yet, caller will retry.
|
||||
return;
|
||||
}
|
||||
|
||||
size_t bytes_remaining =
|
||||
pending_send.pCommand->mSize - pending_send.SizeCurrent;
|
||||
if (bytes_remaining == 0) return;
|
||||
|
||||
EMSCRIPTEN_RESULT result = emscripten_websocket_send_binary(
|
||||
client_socket->mSocket,
|
||||
reinterpret_cast<uint8_t*>(pending_send.pCommand) +
|
||||
pending_send.SizeCurrent,
|
||||
bytes_remaining);
|
||||
|
||||
if (result == EMSCRIPTEN_RESULT_SUCCESS) {
|
||||
pending_send.SizeCurrent += bytes_remaining;
|
||||
pending_send.bError = false;
|
||||
} else {
|
||||
pending_send.bError = true;
|
||||
}
|
||||
}
|
||||
|
||||
SocketInfo* ListenStart(uint32_t /*listen_port*/) {
|
||||
return nullptr; // Browsers cannot open listening ports.
|
||||
}
|
||||
|
||||
SocketInfo* ListenConnect(SocketInfo* /*listen_socket*/) { return nullptr; }
|
||||
|
||||
int GetCloseCode(SocketInfo* client_socket) {
|
||||
return client_socket ? client_socket->mCloseCode.load() : 0;
|
||||
}
|
||||
|
||||
ReadyState GetReadyState(SocketInfo* client_socket) {
|
||||
if (!client_socket) return ReadyState::kDisconnected;
|
||||
if (client_socket->mError) return ReadyState::kError;
|
||||
if (client_socket->mClosed) return ReadyState::kClosed;
|
||||
|
||||
uint16_t ready_state = 0;
|
||||
emscripten_websocket_get_ready_state(client_socket->mSocket, &ready_state);
|
||||
switch (ready_state) {
|
||||
case 0:
|
||||
return ReadyState::kConnecting;
|
||||
case 1:
|
||||
return ReadyState::kOpen;
|
||||
case 2:
|
||||
return ReadyState::kClosing;
|
||||
case 3:
|
||||
return ReadyState::kClosed;
|
||||
}
|
||||
return ReadyState::kError;
|
||||
}
|
||||
|
||||
} // namespace Network
|
||||
} // namespace Internal
|
||||
} // namespace NetImgui
|
||||
Reference in New Issue
Block a user