libs/asio/example/cpp11/allocation/server.cpp

   1 //
   2 // server.cpp
   3 // ~~~~~~~~~~
   4 //
   5 // Copyright (c) 2003-2014 Christopher M. Kohlhoff (chris at kohlhoff dot com)
   6 //
   7 // Distributed under the Boost Software License, Version 1.0. (See accompanying
   8 // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
   9 //
  10
  11 #include <array>
  12 #include <cstdlib>
  13 #include <iostream>
  14 #include <memory>
  15 #include <type_traits>
  16 #include <utility>
  17 #include <boost/asio.hpp>
  18
  19 using boost::asio::ip::tcp;
  20
  21 // Class to manage the memory to be used for handler-based custom allocation.
  22 // It contains a single block of memory which may be returned for allocation
  23 // requests. If the memory is in use when an allocation request is made, the
  24 // allocator delegates allocation to the global heap.
  25 class handler_allocator
  26 {
  27 public:
  28   handler_allocator()
  29     : in_use_(false)
  30   {
  31   }
  32
  33   handler_allocator(const handler_allocator&) = delete;
  34   handler_allocator& operator=(const handler_allocator&) = delete;
  35
  36   void* allocate(std::size_t size)
  37   {
  38     if (!in_use_ && size < sizeof(storage_))
  39     {
  40       in_use_ = true;
  41       return &storage_;
  42     }
  43     else
  44     {
  45       return ::operator new(size);
  46     }
  47   }
  48
  49   void deallocate(void* pointer)
  50   {
  51     if (pointer == &storage_)
  52     {
  53       in_use_ = false;
  54     }
  55     else
  56     {
  57       ::operator delete(pointer);
  58     }
  59   }
  60
  61 private:
  62   // Storage space used for handler-based custom memory allocation.
  63   typename std::aligned_storage<1024>::type storage_;
  64
  65   // Whether the handler-based custom allocation storage has been used.
  66   bool in_use_;
  67 };
  68
  69 // Wrapper class template for handler objects to allow handler memory
  70 // allocation to be customised. Calls to operator() are forwarded to the
  71 // encapsulated handler.
  72 template <typename Handler>
  73 class custom_alloc_handler
  74 {
  75 public:
  76   custom_alloc_handler(handler_allocator& a, Handler h)
  77     : allocator_(a),
  78       handler_(h)
  79   {
  80   }
  81
  82   template <typename ...Args>
  83   void operator()(Args&&... args)
  84   {
  85     handler_(std::forward<Args>(args)...);
  86   }
  87
  88   friend void* asio_handler_allocate(std::size_t size,
  89       custom_alloc_handler<Handler>* this_handler)
  90   {
  91     return this_handler->allocator_.allocate(size);
  92   }
  93
  94   friend void asio_handler_deallocate(void* pointer, std::size_t /*size*/,
  95       custom_alloc_handler<Handler>* this_handler)
  96   {
  97     this_handler->allocator_.deallocate(pointer);
  98   }
  99
 100 private:
 101   handler_allocator& allocator_;
 102   Handler handler_;
 103 };
 104
 105 // Helper function to wrap a handler object to add custom allocation.
 106 template <typename Handler>
 107 inline custom_alloc_handler<Handler> make_custom_alloc_handler(
 108     handler_allocator& a, Handler h)
 109 {
 110   return custom_alloc_handler<Handler>(a, h);
 111 }
 112
 113 class session
 114   : public std::enable_shared_from_this<session>
 115 {
 116 public:
 117   session(tcp::socket socket)
 118     : socket_(std::move(socket))
 119   {
 120   }
 121
 122   void start()
 123   {
 124     do_read();
 125   }
 126
 127 private:
 128   void do_read()
 129   {
 130     auto self(shared_from_this());
 131     socket_.async_read_some(boost::asio::buffer(data_),
 132         make_custom_alloc_handler(allocator_,
 133           [this, self](boost::system::error_code ec, std::size_t length)
 134           {
 135             if (!ec)
 136             {
 137               do_write(length);
 138             }
 139           }));
 140   }
 141
 142   void do_write(std::size_t length)
 143   {
 144     auto self(shared_from_this());
 145     boost::asio::async_write(socket_, boost::asio::buffer(data_, length),
 146         make_custom_alloc_handler(allocator_,
 147           [this, self](boost::system::error_code ec, std::size_t /*length*/)
 148           {
 149             if (!ec)
 150             {
 151               do_read();
 152             }
 153           }));
 154   }
 155
 156   // The socket used to communicate with the client.
 157   tcp::socket socket_;
 158
 159   // Buffer used to store data received from the client.
 160   std::array<char, 1024> data_;
 161
 162   // The allocator to use for handler-based custom memory allocation.
 163   handler_allocator allocator_;
 164 };
 165
 166 class server
 167 {
 168 public:
 169   server(boost::asio::io_service& io_service, short port)
 170     : acceptor_(io_service, tcp::endpoint(tcp::v4(), port)),
 171       socket_(io_service)
 172   {
 173     do_accept();
 174   }
 175
 176 private:
 177   void do_accept()
 178   {
 179     acceptor_.async_accept(socket_,
 180         [this](boost::system::error_code ec)
 181         {
 182           if (!ec)
 183           {
 184             std::make_shared<session>(std::move(socket_))->start();
 185           }
 186
 187           do_accept();
 188         });
 189   }
 190
 191   tcp::acceptor acceptor_;
 192   tcp::socket socket_;
 193 };
 194
 195 int main(int argc, char* argv[])
 196 {
 197   try
 198   {
 199     if (argc != 2)
 200     {
 201       std::cerr << "Usage: server <port>\n";
 202       return 1;
 203     }
 204
 205     boost::asio::io_service io_service;
 206     server s(io_service, std::atoi(argv[1]));
 207     io_service.run();
 208   }
 209   catch (std::exception& e)
 210   {
 211     std::cerr << "Exception: " << e.what() << "\n";
 212   }
 213
 214   return 0;
 215 }