libs/asio/test/latency/udp_server.cpp

   1 //
   2 // udp_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 <boost/asio/io_service.hpp>
  12 #include <boost/asio/ip/udp.hpp>
  13 #include <boost/shared_ptr.hpp>
  14 #include <cstdio>
  15 #include <cstdlib>
  16 #include <cstring>
  17 #include <vector>
  18 #include "allocator.hpp"
  19
  20 using boost::asio::ip::udp;
  21
  22 #include "yield.hpp"
  23
  24 class udp_server : coroutine
  25 {
  26 public:
  27   udp_server(boost::asio::io_service& io_service,
  28       unsigned short port, std::size_t buf_size) :
  29     socket_(io_service, udp::endpoint(udp::v4(), port)),
  30     buffer_(buf_size)
  31   {
  32   }
  33
  34   void operator()(boost::system::error_code ec, std::size_t n = 0)
  35   {
  36     reenter (this) for (;;)
  37     {
  38       yield socket_.async_receive_from(
  39           boost::asio::buffer(buffer_),
  40           sender_, ref(this));
  41
  42       if (!ec)
  43       {
  44         for (std::size_t i = 0; i < n; ++i) buffer_[i] = ~buffer_[i];
  45         socket_.send_to(boost::asio::buffer(buffer_, n), sender_, 0, ec);
  46       }
  47     }
  48   }
  49
  50   friend void* asio_handler_allocate(std::size_t n, udp_server* s)
  51   {
  52     return s->allocator_.allocate(n);
  53   }
  54
  55   friend void asio_handler_deallocate(void* p, std::size_t, udp_server* s)
  56   {
  57     s->allocator_.deallocate(p);
  58   }
  59
  60   struct ref
  61   {
  62     explicit ref(udp_server* p)
  63       : p_(p)
  64     {
  65     }
  66
  67     void operator()(boost::system::error_code ec, std::size_t n = 0)
  68     {
  69       (*p_)(ec, n);
  70     }
  71
  72   private:
  73     udp_server* p_;
  74
  75     friend void* asio_handler_allocate(std::size_t n, ref* r)
  76     {
  77       return asio_handler_allocate(n, r->p_);
  78     }
  79
  80     friend void asio_handler_deallocate(void* p, std::size_t n, ref* r)
  81     {
  82       asio_handler_deallocate(p, n, r->p_);
  83     }
  84   };
  85
  86 private:
  87   udp::socket socket_;
  88   std::vector<unsigned char> buffer_;
  89   udp::endpoint sender_;
  90   allocator allocator_;
  91 };
  92
  93 #include "unyield.hpp"
  94
  95 int main(int argc, char* argv[])
  96 {
  97   if (argc != 5)
  98   {
  99     std::fprintf(stderr,
 100         "Usage: udp_server <port1> <nports> "
 101         "<bufsize> {spin|block}\n");
 102     return 1;
 103   }
 104
 105   unsigned short first_port = static_cast<unsigned short>(std::atoi(argv[1]));
 106   unsigned short num_ports = static_cast<unsigned short>(std::atoi(argv[2]));
 107   std::size_t buf_size = std::atoi(argv[3]);
 108   bool spin = (std::strcmp(argv[4], "spin") == 0);
 109
 110   boost::asio::io_service io_service(1);
 111   std::vector<boost::shared_ptr<udp_server> > servers;
 112
 113   for (unsigned short i = 0; i < num_ports; ++i)
 114   {
 115     unsigned short port = first_port + i;
 116     boost::shared_ptr<udp_server> s(new udp_server(io_service, port, buf_size));
 117     servers.push_back(s);
 118     (*s)(boost::system::error_code());
 119   }
 120
 121   if (spin)
 122     for (;;) io_service.poll();
 123   else
 124     io_service.run();
 125 }