thrift_asio_server.hpp

//
// Created by Benjamin Schulz on 15/03/15.
//

#ifndef _THRIFT_ASIO_SERVER_HPP_
#define _THRIFT_ASIO_SERVER_HPP_

#pragma once

#include <boost/smart_ptr.hpp>
#include <boost/asio.hpp>
#include <thrift/TProcessor.h>
#include <thrift/transport/TBufferTransports.h>
#include <thrift/protocol/TBinaryProtocol.h>
#include <iostream>
#include "./thrift_asio_transport.hpp"

namespace betabugs{
namespace networking{

template <typename HandlerType>
class thrift_asio_server
{
    typedef boost::shared_ptr<HandlerType> Handler_ptr;

    // forward typedefs to minimize pollution
    typedef apache::thrift::TProcessor TProcessor;
    typedef apache::thrift::transport::TMemoryBuffer TMemoryBuffer;
    typedef apache::thrift::protocol::TBinaryProtocol TBinaryProtocol;

  public:
    static void serve(
        boost::asio::io_service& io_service,
        TProcessor& processor,
        Handler_ptr handler,
        unsigned short port
    )
    {
        using boost::asio::ip::tcp;

        auto acceptor = std::make_shared<tcp::acceptor>(
            io_service,
            boost::asio::ip::tcp::endpoint(boost::asio::ip::tcp::v4(), port),
            true
        );

        start_accept(io_service, acceptor, processor, handler);
    }

  private:
    static void start_accept(
        boost::asio::io_service& io_service,
        std::shared_ptr<boost::asio::ip::tcp::acceptor> acceptor,
        TProcessor& processor,
        Handler_ptr handler
    )
    {
        using boost::asio::ip::tcp;

        auto socket = std::make_shared<tcp::socket>(io_service);
        acceptor->async_accept(
            *socket,
            [&io_service, acceptor, socket, &processor, handler]
                (boost::system::error_code ec)
            {
                if (ec)
                {
                    std::clog << ec.message() << std::endl;
                }
                else
                {
                    std::clog << "client connected" << std::endl;
                    on_accept(io_service, socket, processor, handler);
                }

                // Note: this will accept new connections without any bounds
                start_accept(io_service, acceptor, processor, handler);
            }
        );
    }

    // called when a new client connection was established (accepted)
    static void on_accept(
        boost::asio::io_service& io_service,
        std::shared_ptr<boost::asio::ip::tcp::socket> socket,
        TProcessor& processor,
        Handler_ptr handler
    )
    {
        using boost::make_shared;

        // construct the output_protocol and call the handler
        auto t1 = boost::make_shared<thrift_asio_transport>(socket, handler.get());
        auto t2 = boost::make_shared<apache::thrift::transport::TFramedTransport>(t1);
        auto output_protocol = boost::make_shared<TBinaryProtocol>(t2);
        handler->on_client_connected(output_protocol);

        read_frame_size(io_service, socket, output_protocol, processor, handler);
    }

    // read the size of a frame. Clients are expected to use the framed protocol
    static void read_frame_size(
        boost::asio::io_service& io_service,
        std::shared_ptr<boost::asio::ip::tcp::socket> socket,
        boost::shared_ptr<TBinaryProtocol> output_protocol,
        TProcessor& processor,
        Handler_ptr handler
    )
    {
        auto frame_size = std::make_shared<uint32_t>(0);
        boost::asio::async_read(
            *socket, boost::asio::buffer(frame_size.get(), sizeof(uint32_t)),
            [&io_service, socket, output_protocol, &processor, handler, frame_size]
                (const boost::system::error_code& ec, std::size_t /*bytes_transferred*/)
            {
                if(ec)
                {
                    std::clog << ec.message() << std::endl;
                    handler->on_client_disconnected(output_protocol, ec);
                }
                else
                {
                    *frame_size = ntohl(*frame_size);
                    read_frame_data(io_service, socket, output_protocol, processor, handler, *frame_size);
                }
            }
        );
    }

    // read the data of the frame
    static void read_frame_data(
        boost::asio::io_service& io_service,
        std::shared_ptr<boost::asio::ip::tcp::socket> socket,
        boost::shared_ptr<TBinaryProtocol> output_protocol,
        TProcessor& processor,
        Handler_ptr handler,
        uint32_t frame_size
    )
    {
        auto frame_bytes = std::make_shared<std::vector<uint8_t>>(frame_size);
        boost::asio::async_read(
            *socket, boost::asio::buffer(frame_bytes->data(), frame_size),
            [&io_service, socket, output_protocol, &processor, handler, frame_bytes]
                (const boost::system::error_code& ec, std::size_t /*bytes_transferred*/)
            {
                if(ec)
                {
                    std::clog << ec.message() << std::endl;
                    handler->on_client_disconnected(output_protocol, ec);
                }
                else
                {
                    std::clog << "got frame bytes: " << frame_bytes->size() << std::endl;

                    auto input_transport = boost::make_shared<TMemoryBuffer>(frame_bytes->data(), frame_bytes->size());
                    auto input_protocol = boost::make_shared<TBinaryProtocol>(input_transport);

                    void* connection_context = nullptr;

                    handler->before_process(output_protocol);
                    processor.process(input_protocol, output_protocol, connection_context);
                    handler->after_process();

                    // read the next frame
                    read_frame_size(io_service, socket, output_protocol, processor, handler);
                }
            }
        );
    }
};

}
}


#endif //_THRIFT_ASIO_SERVER_HPP_