mirror of
https://github.com/XRPLF/clio.git
synced 2026-07-30 18:40:24 +00:00
cleanup websocket sessions that are subscribed to books or accounts (#146)
This commit is contained in:
@@ -5,25 +5,24 @@
|
||||
template <class T>
|
||||
inline void
|
||||
sendToSubscribers(
|
||||
std::shared_ptr<Message>& message,
|
||||
std::shared_ptr<Message> const& message,
|
||||
T& subscribers,
|
||||
boost::asio::io_context::strand& strand)
|
||||
std::atomic_uint64_t& counter)
|
||||
{
|
||||
boost::asio::post(strand, [&subscribers, message]() {
|
||||
for (auto it = subscribers.begin(); it != subscribers.end();)
|
||||
for (auto it = subscribers.begin(); it != subscribers.end();)
|
||||
{
|
||||
auto& session = *it;
|
||||
if (session->dead())
|
||||
{
|
||||
auto& session = *it;
|
||||
if (session->dead())
|
||||
{
|
||||
it = subscribers.erase(it);
|
||||
}
|
||||
else
|
||||
{
|
||||
session->send(message);
|
||||
++it;
|
||||
}
|
||||
it = subscribers.erase(it);
|
||||
--counter;
|
||||
}
|
||||
});
|
||||
else
|
||||
{
|
||||
session->send(message);
|
||||
++it;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <class T>
|
||||
@@ -31,11 +30,13 @@ inline void
|
||||
addSession(
|
||||
std::shared_ptr<WsBase> session,
|
||||
T& subscribers,
|
||||
boost::asio::io_context::strand& strand)
|
||||
std::atomic_uint64_t& counter)
|
||||
{
|
||||
boost::asio::post(strand, [&subscribers, s = std::move(session)]() {
|
||||
subscribers.emplace(s);
|
||||
});
|
||||
if (!subscribers.contains(session))
|
||||
{
|
||||
subscribers.insert(session);
|
||||
++counter;
|
||||
}
|
||||
}
|
||||
|
||||
template <class T>
|
||||
@@ -43,29 +44,37 @@ inline void
|
||||
removeSession(
|
||||
std::shared_ptr<WsBase> session,
|
||||
T& subscribers,
|
||||
boost::asio::io_context::strand& strand)
|
||||
std::atomic_uint64_t& counter)
|
||||
{
|
||||
boost::asio::post(strand, [&subscribers, s = std::move(session)]() {
|
||||
subscribers.erase(s);
|
||||
});
|
||||
if (subscribers.contains(session))
|
||||
{
|
||||
subscribers.erase(session);
|
||||
--counter;
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
Subscription::subscribe(std::shared_ptr<WsBase> const& session)
|
||||
{
|
||||
addSession(session, subscribers_, strand_);
|
||||
boost::asio::post(strand_, [this, session]() {
|
||||
addSession(session, subscribers_, subCount_);
|
||||
});
|
||||
}
|
||||
|
||||
void
|
||||
Subscription::unsubscribe(std::shared_ptr<WsBase> const& session)
|
||||
{
|
||||
removeSession(session, subscribers_, strand_);
|
||||
boost::asio::post(strand_, [this, session]() {
|
||||
removeSession(session, subscribers_, subCount_);
|
||||
});
|
||||
}
|
||||
|
||||
void
|
||||
Subscription::publish(std::shared_ptr<Message>& message)
|
||||
{
|
||||
sendToSubscribers(message, subscribers_, strand_);
|
||||
boost::asio::post(strand_, [this, message]() {
|
||||
sendToSubscribers(message, subscribers_, subCount_);
|
||||
});
|
||||
}
|
||||
|
||||
template <class Key>
|
||||
@@ -74,7 +83,9 @@ SubscriptionMap<Key>::subscribe(
|
||||
std::shared_ptr<WsBase> const& session,
|
||||
Key const& account)
|
||||
{
|
||||
addSession(session, subscribers_[account], strand_);
|
||||
boost::asio::post(strand_, [this, session, account]() {
|
||||
addSession(session, subscribers_[account], subCount_);
|
||||
});
|
||||
}
|
||||
|
||||
template <class Key>
|
||||
@@ -83,7 +94,22 @@ SubscriptionMap<Key>::unsubscribe(
|
||||
std::shared_ptr<WsBase> const& session,
|
||||
Key const& account)
|
||||
{
|
||||
removeSession(session, subscribers_[account], strand_);
|
||||
boost::asio::post(strand_, [this, account, session]() {
|
||||
if (!subscribers_.contains(account))
|
||||
return;
|
||||
|
||||
if (!subscribers_[account].contains(session))
|
||||
return;
|
||||
|
||||
--subCount_;
|
||||
|
||||
subscribers_[account].erase(session);
|
||||
|
||||
if (subscribers_[account].size() == 0)
|
||||
{
|
||||
subscribers_.erase(account);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <class Key>
|
||||
@@ -92,7 +118,12 @@ SubscriptionMap<Key>::publish(
|
||||
std::shared_ptr<Message>& message,
|
||||
Key const& account)
|
||||
{
|
||||
sendToSubscribers(message, subscribers_[account], strand_);
|
||||
boost::asio::post(strand_, [this, account, message]() {
|
||||
if (!subscribers_.contains(account))
|
||||
return;
|
||||
|
||||
sendToSubscribers(message, subscribers_[account], subCount_);
|
||||
});
|
||||
}
|
||||
|
||||
boost::json::object
|
||||
@@ -122,7 +153,7 @@ getLedgerPubMessage(
|
||||
boost::json::object
|
||||
SubscriptionManager::subLedger(
|
||||
boost::asio::yield_context& yield,
|
||||
std::shared_ptr<WsBase>& session)
|
||||
std::shared_ptr<WsBase> session)
|
||||
{
|
||||
ledgerSubscribers_.subscribe(session);
|
||||
|
||||
@@ -146,19 +177,19 @@ SubscriptionManager::subLedger(
|
||||
}
|
||||
|
||||
void
|
||||
SubscriptionManager::unsubLedger(std::shared_ptr<WsBase>& session)
|
||||
SubscriptionManager::unsubLedger(std::shared_ptr<WsBase> session)
|
||||
{
|
||||
ledgerSubscribers_.unsubscribe(session);
|
||||
}
|
||||
|
||||
void
|
||||
SubscriptionManager::subTransactions(std::shared_ptr<WsBase>& session)
|
||||
SubscriptionManager::subTransactions(std::shared_ptr<WsBase> session)
|
||||
{
|
||||
txSubscribers_.subscribe(session);
|
||||
}
|
||||
|
||||
void
|
||||
SubscriptionManager::unsubTransactions(std::shared_ptr<WsBase>& session)
|
||||
SubscriptionManager::unsubTransactions(std::shared_ptr<WsBase> session)
|
||||
{
|
||||
txSubscribers_.unsubscribe(session);
|
||||
}
|
||||
@@ -169,6 +200,11 @@ SubscriptionManager::subAccount(
|
||||
std::shared_ptr<WsBase>& session)
|
||||
{
|
||||
accountSubscribers_.subscribe(session, account);
|
||||
|
||||
std::unique_lock lk(cleanupMtx_);
|
||||
cleanupFuncs_[session].emplace_back([this, account](session_ptr session) {
|
||||
unsubAccount(account, session);
|
||||
});
|
||||
}
|
||||
|
||||
void
|
||||
@@ -182,15 +218,19 @@ SubscriptionManager::unsubAccount(
|
||||
void
|
||||
SubscriptionManager::subBook(
|
||||
ripple::Book const& book,
|
||||
std::shared_ptr<WsBase>& session)
|
||||
std::shared_ptr<WsBase> session)
|
||||
{
|
||||
bookSubscribers_.subscribe(session, book);
|
||||
|
||||
std::unique_lock lk(cleanupMtx_);
|
||||
cleanupFuncs_[session].emplace_back(
|
||||
[this, book](session_ptr session) { unsubBook(book, session); });
|
||||
}
|
||||
|
||||
void
|
||||
SubscriptionManager::unsubBook(
|
||||
ripple::Book const& book,
|
||||
std::shared_ptr<WsBase>& session)
|
||||
std::shared_ptr<WsBase> session)
|
||||
{
|
||||
bookSubscribers_.unsubscribe(session, book);
|
||||
}
|
||||
@@ -335,31 +375,31 @@ SubscriptionManager::forwardValidation(boost::json::object const& response)
|
||||
void
|
||||
SubscriptionManager::subProposedAccount(
|
||||
ripple::AccountID const& account,
|
||||
std::shared_ptr<WsBase>& session)
|
||||
std::shared_ptr<WsBase> session)
|
||||
{
|
||||
accountProposedSubscribers_.subscribe(session, account);
|
||||
}
|
||||
|
||||
void
|
||||
SubscriptionManager::subManifest(std::shared_ptr<WsBase>& session)
|
||||
SubscriptionManager::subManifest(std::shared_ptr<WsBase> session)
|
||||
{
|
||||
manifestSubscribers_.subscribe(session);
|
||||
}
|
||||
|
||||
void
|
||||
SubscriptionManager::unsubManifest(std::shared_ptr<WsBase>& session)
|
||||
SubscriptionManager::unsubManifest(std::shared_ptr<WsBase> session)
|
||||
{
|
||||
manifestSubscribers_.unsubscribe(session);
|
||||
}
|
||||
|
||||
void
|
||||
SubscriptionManager::subValidation(std::shared_ptr<WsBase>& session)
|
||||
SubscriptionManager::subValidation(std::shared_ptr<WsBase> session)
|
||||
{
|
||||
validationsSubscribers_.subscribe(session);
|
||||
}
|
||||
|
||||
void
|
||||
SubscriptionManager::unsubValidation(std::shared_ptr<WsBase>& session)
|
||||
SubscriptionManager::unsubValidation(std::shared_ptr<WsBase> session)
|
||||
{
|
||||
validationsSubscribers_.unsubscribe(session);
|
||||
}
|
||||
@@ -367,19 +407,34 @@ SubscriptionManager::unsubValidation(std::shared_ptr<WsBase>& session)
|
||||
void
|
||||
SubscriptionManager::unsubProposedAccount(
|
||||
ripple::AccountID const& account,
|
||||
std::shared_ptr<WsBase>& session)
|
||||
std::shared_ptr<WsBase> session)
|
||||
{
|
||||
accountProposedSubscribers_.unsubscribe(session, account);
|
||||
}
|
||||
|
||||
void
|
||||
SubscriptionManager::subProposedTransactions(std::shared_ptr<WsBase>& session)
|
||||
SubscriptionManager::subProposedTransactions(std::shared_ptr<WsBase> session)
|
||||
{
|
||||
txProposedSubscribers_.subscribe(session);
|
||||
}
|
||||
|
||||
void
|
||||
SubscriptionManager::unsubProposedTransactions(std::shared_ptr<WsBase>& session)
|
||||
SubscriptionManager::unsubProposedTransactions(std::shared_ptr<WsBase> session)
|
||||
{
|
||||
txProposedSubscribers_.unsubscribe(session);
|
||||
}
|
||||
|
||||
void
|
||||
SubscriptionManager::cleanup(std::shared_ptr<WsBase> session)
|
||||
{
|
||||
std::unique_lock lk(cleanupMtx_);
|
||||
if (!cleanupFuncs_.contains(session))
|
||||
return;
|
||||
|
||||
for (auto f : cleanupFuncs_[session])
|
||||
{
|
||||
f(session);
|
||||
}
|
||||
|
||||
cleanupFuncs_.erase(session);
|
||||
}
|
||||
|
||||
@@ -11,6 +11,7 @@ class Subscription
|
||||
{
|
||||
boost::asio::io_context::strand strand_;
|
||||
std::unordered_set<std::shared_ptr<WsBase>> subscribers_ = {};
|
||||
std::atomic_uint64_t subCount_ = 0;
|
||||
|
||||
public:
|
||||
Subscription() = delete;
|
||||
@@ -31,15 +32,23 @@ public:
|
||||
|
||||
void
|
||||
publish(std::shared_ptr<Message>& message);
|
||||
|
||||
std::uint64_t
|
||||
count()
|
||||
{
|
||||
return subCount_.load();
|
||||
}
|
||||
};
|
||||
|
||||
template <class Key>
|
||||
class SubscriptionMap
|
||||
{
|
||||
using subscribers = std::unordered_set<std::shared_ptr<WsBase>>;
|
||||
using ptr = std::shared_ptr<WsBase>;
|
||||
using subscribers = std::set<ptr>;
|
||||
|
||||
boost::asio::io_context::strand strand_;
|
||||
std::unordered_map<Key, subscribers> subscribers_ = {};
|
||||
std::atomic_uint64_t subCount_ = 0;
|
||||
|
||||
public:
|
||||
SubscriptionMap() = delete;
|
||||
@@ -60,10 +69,18 @@ public:
|
||||
|
||||
void
|
||||
publish(std::shared_ptr<Message>& message, Key const& key);
|
||||
|
||||
std::uint64_t
|
||||
count()
|
||||
{
|
||||
return subCount_.load();
|
||||
}
|
||||
};
|
||||
|
||||
class SubscriptionManager
|
||||
{
|
||||
using session_ptr = std::shared_ptr<WsBase>;
|
||||
|
||||
std::vector<std::thread> workers_;
|
||||
boost::asio::io_context ioc_;
|
||||
std::optional<boost::asio::io_context::work> work_;
|
||||
@@ -133,9 +150,7 @@ public:
|
||||
}
|
||||
|
||||
boost::json::object
|
||||
subLedger(
|
||||
boost::asio::yield_context& yield,
|
||||
std::shared_ptr<WsBase>& session);
|
||||
subLedger(boost::asio::yield_context& yield, session_ptr session);
|
||||
|
||||
void
|
||||
pubLedger(
|
||||
@@ -145,13 +160,13 @@ public:
|
||||
std::uint32_t txnCount);
|
||||
|
||||
void
|
||||
unsubLedger(std::shared_ptr<WsBase>& session);
|
||||
unsubLedger(session_ptr session);
|
||||
|
||||
void
|
||||
subTransactions(std::shared_ptr<WsBase>& session);
|
||||
subTransactions(session_ptr session);
|
||||
|
||||
void
|
||||
unsubTransactions(std::shared_ptr<WsBase>& session);
|
||||
unsubTransactions(session_ptr session);
|
||||
|
||||
void
|
||||
pubTransaction(
|
||||
@@ -159,32 +174,28 @@ public:
|
||||
ripple::LedgerInfo const& lgrInfo);
|
||||
|
||||
void
|
||||
subAccount(
|
||||
ripple::AccountID const& account,
|
||||
std::shared_ptr<WsBase>& session);
|
||||
subAccount(ripple::AccountID const& account, session_ptr& session);
|
||||
|
||||
void
|
||||
unsubAccount(
|
||||
ripple::AccountID const& account,
|
||||
std::shared_ptr<WsBase>& session);
|
||||
unsubAccount(ripple::AccountID const& account, session_ptr& session);
|
||||
|
||||
void
|
||||
subBook(ripple::Book const& book, std::shared_ptr<WsBase>& session);
|
||||
subBook(ripple::Book const& book, session_ptr session);
|
||||
|
||||
void
|
||||
unsubBook(ripple::Book const& book, std::shared_ptr<WsBase>& session);
|
||||
unsubBook(ripple::Book const& book, session_ptr session);
|
||||
|
||||
void
|
||||
subManifest(std::shared_ptr<WsBase>& session);
|
||||
subManifest(session_ptr session);
|
||||
|
||||
void
|
||||
unsubManifest(std::shared_ptr<WsBase>& session);
|
||||
unsubManifest(session_ptr session);
|
||||
|
||||
void
|
||||
subValidation(std::shared_ptr<WsBase>& session);
|
||||
subValidation(session_ptr session);
|
||||
|
||||
void
|
||||
unsubValidation(std::shared_ptr<WsBase>& session);
|
||||
unsubValidation(session_ptr session);
|
||||
|
||||
void
|
||||
forwardProposedTransaction(boost::json::object const& response);
|
||||
@@ -196,26 +207,51 @@ public:
|
||||
forwardValidation(boost::json::object const& response);
|
||||
|
||||
void
|
||||
subProposedAccount(
|
||||
ripple::AccountID const& account,
|
||||
std::shared_ptr<WsBase>& session);
|
||||
subProposedAccount(ripple::AccountID const& account, session_ptr session);
|
||||
|
||||
void
|
||||
unsubProposedAccount(
|
||||
ripple::AccountID const& account,
|
||||
std::shared_ptr<WsBase>& session);
|
||||
unsubProposedAccount(ripple::AccountID const& account, session_ptr session);
|
||||
|
||||
void
|
||||
subProposedTransactions(std::shared_ptr<WsBase>& session);
|
||||
subProposedTransactions(session_ptr session);
|
||||
|
||||
void
|
||||
unsubProposedTransactions(std::shared_ptr<WsBase>& session);
|
||||
unsubProposedTransactions(session_ptr session);
|
||||
|
||||
void
|
||||
cleanup(session_ptr session);
|
||||
|
||||
boost::json::object
|
||||
report()
|
||||
{
|
||||
boost::json::object counts = {};
|
||||
|
||||
counts["ledger"] = ledgerSubscribers_.count();
|
||||
counts["transactions"] = txSubscribers_.count();
|
||||
counts["transactions_proposed"] = txProposedSubscribers_.count();
|
||||
counts["manifests"] = manifestSubscribers_.count();
|
||||
counts["validations"] = validationsSubscribers_.count();
|
||||
counts["account"] = accountSubscribers_.count();
|
||||
counts["accounts_proposed"] = accountProposedSubscribers_.count();
|
||||
counts["books"] = bookSubscribers_.count();
|
||||
|
||||
return counts;
|
||||
}
|
||||
|
||||
private:
|
||||
void
|
||||
sendAll(
|
||||
std::string const& pubMsg,
|
||||
std::unordered_set<std::shared_ptr<WsBase>>& subs);
|
||||
sendAll(std::string const& pubMsg, std::unordered_set<session_ptr>& subs);
|
||||
|
||||
/**
|
||||
* This is how we chose to cleanup subscriptions that have been closed.
|
||||
* Each time we add a subscriber, we add the opposite lambda that
|
||||
* unsubscribes that subscriber when cleanup is called with the session that
|
||||
* closed.
|
||||
*/
|
||||
using CleanupFunction = std::function<void(session_ptr)>;
|
||||
std::mutex cleanupMtx_;
|
||||
std::unordered_map<session_ptr, std::vector<CleanupFunction>>
|
||||
cleanupFuncs_ = {};
|
||||
};
|
||||
|
||||
#endif // SUBSCRIPTION_MANAGER_H
|
||||
|
||||
Reference in New Issue
Block a user