diff --git a/src/QueryScheduler.h b/src/QueryScheduler.h index 72722c82..8a88564b 100644 --- a/src/QueryScheduler.h +++ b/src/QueryScheduler.h @@ -7,6 +7,7 @@ struct QueryScheduler : NonCopyable { std::function onEvent; std::function &levIds)> onEventBatch; std::function onComplete; + std::function onSliceComplete; // If false, then levIds returned to above callbacks can be stale (because they were deleted) // If false, then onEvent's eventPayload will always be "" @@ -97,6 +98,8 @@ struct QueryScheduler : NonCopyable { levIdBatch.clear(); } + if (onSliceComplete) onSliceComplete(txn); + if (complete) { auto connId = q->sub.connId; removeSub(connId, q->sub.subId); diff --git a/src/apps/relay/RelayReqMonitor.cpp b/src/apps/relay/RelayReqMonitor.cpp index 1e6cd4de..78f16fee 100644 --- a/src/apps/relay/RelayReqMonitor.cpp +++ b/src/apps/relay/RelayReqMonitor.cpp @@ -38,8 +38,9 @@ void RelayServer::runReqMonitor(ThreadPool::Thread &thr) { PackedEventView packed(ev.buf); if (msg->sub.filterGroup.doesMatch(packed)) { if (ReadRestrictor::shouldSendToSubscriber(packed, connAuthedPubkey)) { - sendEvent(connId, msg->sub.subId, getEventJson(txn, decomp, ev.primaryKeyId)); + sendEvent(connId, msg->sub.subId, getEventJson(txn, decomp, ev.primaryKeyId), false); } + hubTrigger->send(); } return true; diff --git a/src/apps/relay/RelayReqWorker.cpp b/src/apps/relay/RelayReqWorker.cpp index 56021aea..5d031cc3 100644 --- a/src/apps/relay/RelayReqWorker.cpp +++ b/src/apps/relay/RelayReqWorker.cpp @@ -8,6 +8,9 @@ void RelayServer::runReqWorker(ThreadPool::Thread &thr) { QueryScheduler queries; flat_hash_map connIdToAuthedPubkey; + std::vector pendingPayloads; + uint64_t pendingConnId = 0; + queries.onEvent = [&](lmdb::txn &txn, const auto &sub, uint64_t levId, std::string_view eventPayload){ if (sub.countOnly) return; auto it = connIdToAuthedPubkey.find(sub.connId); @@ -15,10 +18,24 @@ void RelayServer::runReqWorker(ThreadPool::Thread &thr) { PackedEventView packed(ev.buf); Bytes32 subscriberAuthedPubkey = it == connIdToAuthedPubkey.end() ? Bytes32() : it->second; if (!ReadRestrictor::shouldSendToSubscriber(packed, subscriberAuthedPubkey)) { - return; + return; + } + PROM_INC_RELAY_MSG("EVENT"); + pendingConnId = sub.connId; + pendingPayloads.push_back( + buildEventReply(sub.subId, decodeEventPayload(txn, decomp, eventPayload, nullptr, nullptr)) + ); + }; + + queries.onSliceComplete = [&](lmdb::txn &){ + if (pendingPayloads.empty()) return; + if (pendingPayloads.size() == 1) { + tpWebsocket.dispatch(0, MsgWebsocket{MsgWebsocket::Send{pendingConnId, std::move(pendingPayloads[0])}}); + } else { + tpWebsocket.dispatch(0, MsgWebsocket{MsgWebsocket::SendBatch{pendingConnId, std::move(pendingPayloads)}}); } - - sendEvent(sub.connId, sub.subId, decodeEventPayload(txn, decomp, eventPayload, nullptr, nullptr)); + pendingPayloads.clear(); + hubTrigger->send(); }; queries.onComplete = [&](lmdb::txn &, Subscription &sub, uint64_t total){ diff --git a/src/apps/relay/RelayServer.h b/src/apps/relay/RelayServer.h index 7767d8c6..547e1669 100644 --- a/src/apps/relay/RelayServer.h +++ b/src/apps/relay/RelayServer.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -36,6 +37,11 @@ struct MsgWebsocket : NonCopyable { std::string payload; }; + struct SendBatch { + uint64_t connId; + std::vector payloads; + }; + struct SendEventToBatch { RecipientList list; std::string evJson; @@ -44,7 +50,7 @@ struct MsgWebsocket : NonCopyable { struct GracefulShutdown { }; - using Var = std::variant; + using Var = std::variant; Var msg; MsgWebsocket(Var &&msg_) : msg(std::move(msg_)) {} }; @@ -213,65 +219,72 @@ struct RelayServer { // Utils (can be called by any thread) - void sendToConn(uint64_t connId, std::string &&payload) { + void sendToConn(uint64_t connId, std::string &&payload, bool flush = true) { tpWebsocket.dispatch(0, MsgWebsocket{MsgWebsocket::Send{connId, std::move(payload)}}); - hubTrigger->send(); + if (flush) hubTrigger->send(); } - void sendToConnBinary(uint64_t connId, std::string &&payload) { + void sendToConnBinary(uint64_t connId, std::string &&payload, bool flush = true) { tpWebsocket.dispatch(0, MsgWebsocket{MsgWebsocket::SendBinary{connId, std::move(payload)}}); - hubTrigger->send(); + if (flush) hubTrigger->send(); } - void sendEvent(uint64_t connId, const SubId &subId, std::string_view evJson) { - PROM_INC_RELAY_MSG("EVENT"); - auto subIdSv = subId.sv(); + void sendToConnBatched(uint64_t connId, std::vector &&payloads, bool flush = true) { + tpWebsocket.dispatch(0, MsgWebsocket{MsgWebsocket::SendBatch{connId, std::move(payloads)}}); + if (flush) hubTrigger->send(); + } + std::string buildEventReply(const SubId &subId, std::string_view evJson) { + auto subIdSv = subId.sv(); std::string reply; reply.reserve(13 + subIdSv.size() + evJson.size()); - reply += "[\"EVENT\",\""; reply += subIdSv; reply += "\","; reply += evJson; reply += "]"; + return reply; + } - sendToConn(connId, std::move(reply)); + void sendEvent(uint64_t connId, const SubId &subId, std::string_view evJson, bool flush = true) { + PROM_INC_RELAY_MSG("EVENT"); + auto reply = buildEventReply(subId, evJson); + sendToConn(connId, std::move(reply), flush); } - void sendEventToBatch(RecipientList &&list, std::string &&evJson) { + void sendEventToBatch(RecipientList &&list, std::string &&evJson, bool flush = true) { tpWebsocket.dispatch(0, MsgWebsocket{MsgWebsocket::SendEventToBatch{std::move(list), std::move(evJson)}}); - hubTrigger->send(); + if (flush) hubTrigger->send(); } - void sendNoticeError(uint64_t connId, std::string &&payload) { + void sendNoticeError(uint64_t connId, std::string &&payload, bool flush = true) { PROM_INC_RELAY_MSG("NOTICE"); LI << "sending error to [" << connId << "]: " << payload; auto reply = tao::json::value::array({ "NOTICE", std::string("ERROR: ") + payload }); tpWebsocket.dispatch(0, MsgWebsocket{MsgWebsocket::Send{connId, std::move(tao::json::to_string(reply))}}); - hubTrigger->send(); + if (flush) hubTrigger->send(); } - void sendClosedError(uint64_t connId, const std::string &subId, std::string &&payload) { + void sendClosedError(uint64_t connId, const std::string &subId, std::string &&payload, bool flush = true) { PROM_INC_RELAY_MSG("CLOSED"); LI << "sending closed to [" << connId << "]: " << payload; auto reply = tao::json::value::array({ "CLOSED", subId, std::string("ERROR: ") + payload }); tpWebsocket.dispatch(0, MsgWebsocket{MsgWebsocket::Send{connId, std::move(tao::json::to_string(reply))}}); - hubTrigger->send(); + if (flush) hubTrigger->send(); } - void sendOKResponse(uint64_t connId, std::string_view eventIdHex, bool written, std::string_view message) { + void sendOKResponse(uint64_t connId, std::string_view eventIdHex, bool written, std::string_view message, bool flush = true) { PROM_INC_RELAY_MSG("OK"); auto reply = tao::json::value::array({ "OK", eventIdHex, written, message }); tpWebsocket.dispatch(0, MsgWebsocket{MsgWebsocket::Send{connId, std::move(tao::json::to_string(reply))}}); - hubTrigger->send(); + if (flush) hubTrigger->send(); } - void sendAuthChallenge(uint64_t connId, std::string_view challenge) { + void sendAuthChallenge(uint64_t connId, std::string_view challenge, bool flush = true) { PROM_INC_RELAY_MSG("AUTH"); PrometheusMetrics::getInstance().authChallengesSentTotal.inc(); auto reply = tao::json::value::array({ "AUTH", challenge }); tpWebsocket.dispatch(0, MsgWebsocket{MsgWebsocket::Send{connId, std::move(tao::json::to_string(reply))}}); - hubTrigger->send(); + if (flush) hubTrigger->send(); } }; diff --git a/src/apps/relay/RelayWebsocket.cpp b/src/apps/relay/RelayWebsocket.cpp index fbd302a1..a097b7d3 100644 --- a/src/apps/relay/RelayWebsocket.cpp +++ b/src/apps/relay/RelayWebsocket.cpp @@ -331,11 +331,7 @@ void RelayServer::runWebsocket(ThreadPool::Thread &thr) { std::function asyncCb = [&]{ auto newMsgs = thr.inbox.pop_all_no_wait(); - auto doSend = [&](uint64_t connId, std::string_view payload, uWS::OpCode opCode){ - auto it = connIdToConnection.find(connId); - if (it == connIdToConnection.end()) return; - auto &c = *it->second; - + auto doSendToConnection = [&](Connection &c, std::string_view payload, uWS::OpCode opCode) -> bool { // Track bytes still inside uWS's outbound path (either queued or // partially sent). Increment before send(), decrement in the // completion callback. The payload size is smuggled through the @@ -380,8 +376,15 @@ void RelayServer::runWebsocket(ThreadPool::Thread &thr) { << renderSize(maxPending) << ", terminating"; PrometheusMetrics::getInstance().slowClientTerminations.inc(); c.websocket->terminate(); - return; + return false; // signal: connection is dead, stop sending to it } + return true; + }; + + auto doSend = [&](uint64_t connId, std::string_view payload, uWS::OpCode opCode){ + auto it = connIdToConnection.find(connId); + if (it == connIdToConnection.end()) return; + doSendToConnection(*it->second, payload, opCode); }; for (auto &newMsg : newMsgs) { @@ -389,6 +392,13 @@ void RelayServer::runWebsocket(ThreadPool::Thread &thr) { doSend(msg->connId, msg->payload, uWS::OpCode::TEXT); } else if (auto msg = std::get_if(&newMsg.msg)) { doSend(msg->connId, msg->payload, uWS::OpCode::BINARY); + } else if (auto batch = std::get_if(&newMsg.msg)) { + auto it = connIdToConnection.find(batch->connId); + if (it == connIdToConnection.end()) continue; + auto &c = *it->second; + for (auto &payload : batch->payloads) { + if (!doSendToConnection(c, payload, uWS::OpCode::TEXT)) break; // terminated, stop + } } else if (auto msg = std::get_if(&newMsg.msg)) { tempBuf.reserve(13 + MAX_SUBID_SIZE + msg->evJson.size()); tempBuf.resize(10 + MAX_SUBID_SIZE); diff --git a/test/tests/readRestrictTest.js b/test/tests/readRestrictTest.js index c1c3cc0e..05a62e3d 100644 --- a/test/tests/readRestrictTest.js +++ b/test/tests/readRestrictTest.js @@ -210,8 +210,15 @@ async function testRestrictedFilterReturnsAllIfAuthenticatedAndInvolvementNotReq }) { // send AUTH message with challenge string const authEvent = signEvent( - authChallengeString, - "wss://relay.test", + { + kind: 22242, + created_at: Math.floor(Date.now() / 1000), + tags: [ + ["relay", "wss://relay.test"], + ["challenge", authChallengeString], + ], + content: "", + }, ids[0].sec, ); client.send(["AUTH", authEvent]); diff --git a/test/utils/events.js b/test/utils/events.js index e87f3dfb..265e6b6f 100644 --- a/test/utils/events.js +++ b/test/utils/events.js @@ -52,20 +52,7 @@ export function buildEvent({ return { ...evt, id, sig }; } -export function signEvent(authChallengeString, relayUrl, sec) { +export function signEvent(event, sec) { const secBytes = hexToBytes(sec); - const authEvent = finalizeEvent( - { - kind: 22242, - created_at: Math.floor(Date.now() / 1000), - tags: [ - ["relay", relayUrl], - ["challenge", authChallengeString], - ], - content: "", - }, - secBytes, - ); - - return authEvent; + return finalizeEvent(event, secBytes); }