proxy-io.h

Go to the documentation of this file.

// Copyright (c) 2019 The Bitcoin Core developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.
 
#ifndef MP_PROXY_IO_H
#define MP_PROXY_IO_H
 
#include <mp/proxy.h>
#include <mp/util.h>
 
#include <mp/proxy.capnp.h>
 
#include <capnp/rpc-twoparty.h>
 
#include <assert.h>
#include <functional>
#include <optional>
#include <map>
#include <memory>
#include <sstream>
#include <string>
 
namespace mp {
struct ThreadContext;
 
struct InvokeContext
{
    Connection& connection;
};
 
struct ClientInvokeContext : InvokeContext
{
    ThreadContext& thread_context;
    ClientInvokeContext(Connection& conn, ThreadContext& thread_context)
        : InvokeContext{conn}, thread_context{thread_context}
    {
    }
};
 
template <typename ProxyServer, typename CallContext_>
struct ServerInvokeContext : InvokeContext
{
    using CallContext = CallContext_;
 
    ProxyServer& proxy_server;
    CallContext& call_context;
    int req;
 
    ServerInvokeContext(ProxyServer& proxy_server, CallContext& call_context, int req)
        : InvokeContext{*proxy_server.m_context.connection}, proxy_server{proxy_server}, call_context{call_context}, req{req}
    {
    }
};
 
template <typename Interface, typename Params, typename Results>
using ServerContext = ServerInvokeContext<ProxyServer<Interface>, ::capnp::CallContext<Params, Results>>;
 
template <>
struct ProxyClient<Thread> : public ProxyClientBase<Thread, ::capnp::Void>
{
    using ProxyClientBase::ProxyClientBase;
    // https://stackoverflow.com/questions/22357887/comparing-two-mapiterators-why-does-it-need-the-copy-constructor-of-stdpair
    ProxyClient(const ProxyClient&) = delete;
    ~ProxyClient();
 
    void setCleanup(const std::function<void()>& fn);
 
    std::optional<CleanupIt> m_cleanup_it;
};
 
template <>
struct ProxyServer<Thread> final : public Thread::Server
{
public:
    ProxyServer(ThreadContext& thread_context, std::thread&& thread);
    ~ProxyServer();
    kj::Promise<void> getName(GetNameContext context) override;
    ThreadContext& m_thread_context;
    std::thread m_thread;
};
 
class LoggingErrorHandler : public kj::TaskSet::ErrorHandler
{
public:
    LoggingErrorHandler(EventLoop& loop) : m_loop(loop) {}
    void taskFailed(kj::Exception&& exception) override;
    EventLoop& m_loop;
};
 
using LogFn = std::function<void(bool raise, std::string message)>;
 
class Logger
{
public:
    Logger(bool raise, LogFn& fn) : m_raise(raise), m_fn(fn) {}
    Logger(Logger&& logger) : m_raise(logger.m_raise), m_fn(logger.m_fn), m_buffer(std::move(logger.m_buffer)) {}
    ~Logger() noexcept(false)
    {
        if (m_fn) m_fn(m_raise, m_buffer.str());
    }
 
    template <typename T>
    friend Logger& operator<<(Logger& logger, T&& value)
    {
        if (logger.m_fn) logger.m_buffer << std::forward<T>(value);
        return logger;
    }
 
    template <typename T>
    friend Logger& operator<<(Logger&& logger, T&& value)
    {
        return logger << std::forward<T>(value);
    }
 
    bool m_raise;
    LogFn& m_fn;
    std::ostringstream m_buffer;
};
 
std::string LongThreadName(const char* exe_name);
 
class EventLoop
{
public:
    EventLoop(const char* exe_name, LogFn log_fn, void* context = nullptr);
    ~EventLoop();
 
    void loop();
 
    void post(const std::function<void()>& fn);
 
    template <typename Callable>
    void sync(Callable&& callable)
    {
        post(std::ref(callable));
    }
 
    void startAsyncThread(std::unique_lock<std::mutex>& lock);
 
    void addClient(std::unique_lock<std::mutex>& lock);
    bool removeClient(std::unique_lock<std::mutex>& lock);
    bool done(std::unique_lock<std::mutex>& lock);
 
    Logger log()
    {
        Logger logger(false, m_log_fn);
        logger << "{" << LongThreadName(m_exe_name) << "} ";
        return logger;
    }
    Logger logPlain() { return {false, m_log_fn}; }
    Logger raise() { return {true, m_log_fn}; }
 
    const char* m_exe_name;
 
    std::thread::id m_thread_id = std::this_thread::get_id();
 
    std::thread m_async_thread;
 
    const std::function<void()>* m_post_fn = nullptr;
 
    CleanupList m_async_fns;
 
    int m_wait_fd = -1;
 
    int m_post_fd = -1;
 
    int m_num_clients = 0;
 
    std::mutex m_mutex;
    std::condition_variable m_cv;
 
    kj::AsyncIoContext m_io_context;
 
    LoggingErrorHandler m_error_handler{*this};
 
    std::unique_ptr<kj::TaskSet> m_task_set;
 
    std::list<Connection> m_incoming_connections;
 
    LogFn m_log_fn;
 
    void* m_context;
};
 
struct Waiter
{
    Waiter() = default;
 
    template <typename Fn>
    void post(Fn&& fn)
    {
        const std::unique_lock<std::mutex> lock(m_mutex);
        assert(!m_fn);
        m_fn = std::move(fn);
        m_cv.notify_all();
    }
 
    template <class Predicate>
    void wait(std::unique_lock<std::mutex>& lock, Predicate pred)
    {
        m_cv.wait(lock, [&] {
            // Important for this to be "while (m_fn)", not "if (m_fn)" to avoid
            // a lost-wakeup bug. A new m_fn and m_cv notification might be sent
            // after the fn() call and before the lock.lock() call in this loop
            // in the case where a capnp response is sent and a brand new
            // request is immediately received.
            while (m_fn) {
                auto fn = std::move(m_fn);
                m_fn = nullptr;
                lock.unlock();
                fn();
                lock.lock();
            }
            const bool done = pred();
            return done;
        });
    }
 
    std::mutex m_mutex;
    std::condition_variable m_cv;
    std::function<void()> m_fn;
};
 
class Connection
{
public:
    Connection(EventLoop& loop, kj::Own<kj::AsyncIoStream>&& stream_)
        : m_loop(loop), m_stream(kj::mv(stream_)),
          m_network(*m_stream, ::capnp::rpc::twoparty::Side::CLIENT, ::capnp::ReaderOptions()),
          m_rpc_system(::capnp::makeRpcClient(m_network))
    {
        std::unique_lock<std::mutex> lock(m_loop.m_mutex);
        m_loop.addClient(lock);
    }
    Connection(EventLoop& loop,
        kj::Own<kj::AsyncIoStream>&& stream_,
        const std::function<::capnp::Capability::Client(Connection&)>& make_client)
        : m_loop(loop), m_stream(kj::mv(stream_)),
          m_network(*m_stream, ::capnp::rpc::twoparty::Side::SERVER, ::capnp::ReaderOptions()),
          m_rpc_system(::capnp::makeRpcServer(m_network, make_client(*this)))
    {
        std::unique_lock<std::mutex> lock(m_loop.m_mutex);
        m_loop.addClient(lock);
    }
 
    ~Connection();
 
    CleanupIt addSyncCleanup(std::function<void()> fn);
    void removeSyncCleanup(CleanupIt it);
 
    void addAsyncCleanup(std::function<void()> fn);
 
    template <typename F>
    void onDisconnect(F&& f)
    {
        // Add disconnect handler to local TaskSet to ensure it is cancelled and
        // will never run after connection object is destroyed. But when disconnect
        // handler fires, do not call the function f right away, instead add it
        // to the EventLoop TaskSet to avoid "Promise callback destroyed itself"
        // error in cases where f deletes this Connection object.
        m_on_disconnect.add(m_network.onDisconnect().then(
            [f = std::move(f), this]() mutable { m_loop.m_task_set->add(kj::evalLater(kj::mv(f))); }));
    }
 
    EventLoop& m_loop;
    kj::Own<kj::AsyncIoStream> m_stream;
    LoggingErrorHandler m_error_handler{m_loop};
    kj::TaskSet m_on_disconnect{m_error_handler};
    ::capnp::TwoPartyVatNetwork m_network;
    std::optional<::capnp::RpcSystem<::capnp::rpc::twoparty::VatId>> m_rpc_system;
 
    // ThreadMap interface client, used to create a remote server thread when an
    // client IPC call is being made for the first time from a new thread.
    ThreadMap::Client m_thread_map{nullptr};
 
    ::capnp::CapabilityServerSet<Thread> m_threads;
 
    CleanupList m_sync_cleanup_fns;
    CleanupList m_async_cleanup_fns;
};
 
struct ServerVatId
{
    ::capnp::word scratch[4]{};
    ::capnp::MallocMessageBuilder message{scratch};
    ::capnp::rpc::twoparty::VatId::Builder vat_id{message.getRoot<::capnp::rpc::twoparty::VatId>()};
    ServerVatId() { vat_id.setSide(::capnp::rpc::twoparty::Side::SERVER); }
};
 
template <typename Interface, typename Impl>
ProxyClientBase<Interface, Impl>::ProxyClientBase(typename Interface::Client client,
    Connection* connection,
    bool destroy_connection)
    : m_client(std::move(client)), m_context(connection)
 
{
    {
        std::unique_lock<std::mutex> lock(m_context.connection->m_loop.m_mutex);
        m_context.connection->m_loop.addClient(lock);
    }
 
    // Handler for the connection getting destroyed before this client object.
    auto cleanup_it = m_context.connection->addSyncCleanup([this]() {
        // Release client capability by move-assigning to temporary.
        {
            typename Interface::Client(std::move(m_client));
        }
        {
            std::unique_lock<std::mutex> lock(m_context.connection->m_loop.m_mutex);
            m_context.connection->m_loop.removeClient(lock);
        }
        m_context.connection = nullptr;
    });
 
    // Two shutdown sequences are supported:
    //
    // - A normal sequence where client proxy objects are deleted by external
    //   code that no longer needs them
    //
    // - A garbage collection sequence where the connection or event loop shuts
    //   down while external code is still holding client references.
    //
    // The first case is handled here when m_context.connection is not null. The
    // second case is handled by the cleanup function, which sets m_context.connection to
    // null so nothing happens here.
    m_context.cleanup_fns.emplace_front([this, destroy_connection, cleanup_it]{
    if (m_context.connection) {
        // Remove cleanup callback so it doesn't run and try to access
        // this object after it's already destroyed.
        m_context.connection->removeSyncCleanup(cleanup_it);
 
        // If the capnp interface defines a destroy method, call it to destroy
        // the remote object, waiting for it to be deleted server side. If the
        // capnp interface does not define a destroy method, this will just call
        // an empty stub defined in the ProxyClientBase class and do nothing.
        Sub::destroy(*this);
 
        // FIXME: Could just invoke removed addCleanup fn here instead of duplicating code
        m_context.connection->m_loop.sync([&]() {
            // Release client capability by move-assigning to temporary.
            {
                typename Interface::Client(std::move(m_client));
            }
            {
                std::unique_lock<std::mutex> lock(m_context.connection->m_loop.m_mutex);
                m_context.connection->m_loop.removeClient(lock);
            }
 
            if (destroy_connection) {
                delete m_context.connection;
                m_context.connection = nullptr;
            }
        });
    }
    });
    Sub::construct(*this);
}
 
template <typename Interface, typename Impl>
ProxyClientBase<Interface, Impl>::~ProxyClientBase() noexcept
{
    CleanupRun(m_context.cleanup_fns);
}
 
template <typename Interface, typename Impl>
ProxyServerBase<Interface, Impl>::ProxyServerBase(std::shared_ptr<Impl> impl, Connection& connection)
    : m_impl(std::move(impl)), m_context(&connection)
{
    assert(m_impl);
    std::unique_lock<std::mutex> lock(m_context.connection->m_loop.m_mutex);
    m_context.connection->m_loop.addClient(lock);
}
 
template <typename Interface, typename Impl>
ProxyServerBase<Interface, Impl>::~ProxyServerBase()
{
    if (m_impl) {
        // If impl is non-null at this point, it means no client is waiting for
        // the m_impl server object to be destroyed synchronously. This can
        // happen either if the interface did not define a "destroy" method (see
        // invokeDestroy method below), or if a destroy method was defined, but
        // the connection was broken before it could be called.
        //
        // In either case, be conservative and run the cleanup on an
        // asynchronous thread, to avoid destructors or cleanup functions
        // blocking or deadlocking the current EventLoop thread, since they
        // could be making IPC calls.
        //
        // Technically this is a little too conservative since if the interface
        // defines a "destroy" method, but the destroy method does not accept a
        // Context parameter specifying a worker thread, the cleanup method
        // would run on the EventLoop thread normally (when connection is
        // unbroken), but will not run on the EventLoop thread now (when
        // connection is broken). Probably some refactoring of the destructor
        // and invokeDestroy function is possible to make this cleaner and more
        // consistent.
        m_context.connection->addAsyncCleanup([impl=std::move(m_impl), fns=std::move(m_context.cleanup_fns)]() mutable {
            impl.reset();
            CleanupRun(fns);
        });
    }
    assert(m_context.cleanup_fns.empty());
    std::unique_lock<std::mutex> lock(m_context.connection->m_loop.m_mutex);
    m_context.connection->m_loop.removeClient(lock);
}
 
template <typename Interface, typename Impl>
void ProxyServerBase<Interface, Impl>::invokeDestroy()
{
    m_impl.reset();
    CleanupRun(m_context.cleanup_fns);
}
 
using ConnThreads = std::map<Connection*, ProxyClient<Thread>>;
using ConnThread = ConnThreads::iterator;
 
// Retrieve ProxyClient<Thread> object associated with this connection from a
// map, or create a new one and insert it into the map. Return map iterator and
// inserted bool.
std::tuple<ConnThread, bool> SetThread(ConnThreads& threads, std::mutex& mutex, Connection* connection, const std::function<Thread::Client()>& make_thread);
 
struct ThreadContext
{
    std::string thread_name;
 
    std::unique_ptr<Waiter> waiter = nullptr;
 
    ConnThreads callback_threads;
 
    ConnThreads request_threads;
 
    bool loop_thread = false;
};
 
template <typename InitInterface>
std::unique_ptr<ProxyClient<InitInterface>> ConnectStream(EventLoop& loop, int fd)
{
    typename InitInterface::Client init_client(nullptr);
    std::unique_ptr<Connection> connection;
    loop.sync([&] {
        auto stream =
            loop.m_io_context.lowLevelProvider->wrapSocketFd(fd, kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP);
        connection = std::make_unique<Connection>(loop, kj::mv(stream));
        init_client = connection->m_rpc_system->bootstrap(ServerVatId().vat_id).castAs<InitInterface>();
        Connection* connection_ptr = connection.get();
        connection->onDisconnect([&loop, connection_ptr] {
            loop.log() << "IPC client: unexpected network disconnect.";
            delete connection_ptr;
        });
    });
    return std::make_unique<ProxyClient<InitInterface>>(
        kj::mv(init_client), connection.release(), /* destroy_connection= */ true);
}
 
template <typename InitInterface, typename InitImpl>
void _Serve(EventLoop& loop, kj::Own<kj::AsyncIoStream>&& stream, InitImpl& init)
{
    loop.m_incoming_connections.emplace_front(loop, kj::mv(stream), [&](Connection& connection) {
        // Disable deleter so proxy server object doesn't attempt to delete the
        // init implementation when the proxy client is destroyed or
        // disconnected.
        return kj::heap<ProxyServer<InitInterface>>(std::shared_ptr<InitImpl>(&init, [](InitImpl*){}), connection);
    });
    auto it = loop.m_incoming_connections.begin();
    it->onDisconnect([&loop, it] {
        loop.log() << "IPC server: socket disconnected.";
        loop.m_incoming_connections.erase(it);
    });
}
 
template <typename InitInterface, typename InitImpl>
void _Listen(EventLoop& loop, kj::Own<kj::ConnectionReceiver>&& listener, InitImpl& init)
{
    auto* ptr = listener.get();
    loop.m_task_set->add(ptr->accept().then(
        [&loop, &init, listener = kj::mv(listener)](kj::Own<kj::AsyncIoStream>&& stream) mutable {
            _Serve<InitInterface>(loop, kj::mv(stream), init);
            _Listen<InitInterface>(loop, kj::mv(listener), init);
        }));
}
 
template <typename InitInterface, typename InitImpl>
void ServeStream(EventLoop& loop, int fd, InitImpl& init)
{
    _Serve<InitInterface>(
        loop, loop.m_io_context.lowLevelProvider->wrapSocketFd(fd, kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP), init);
}
 
template <typename InitInterface, typename InitImpl>
void ListenConnections(EventLoop& loop, int fd, InitImpl& init)
{
    loop.sync([&]() {
        _Listen<InitInterface>(loop,
            loop.m_io_context.lowLevelProvider->wrapListenSocketFd(fd, kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP),
            init);
    });
}
 
extern thread_local ThreadContext g_thread_context;
 
} // namespace mp
 
#endif // MP_PROXY_IO_H