aboutsummaryrefslogtreecommitdiffhomepage
path: root/tdutils/td/utils/BufferedUdp.h
blob: 4c030afae84417dff0c8320edb0aa67b2f491a05 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
//
// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2025
//
// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//
#pragma once

#include "td/utils/buffer.h"
#include "td/utils/common.h"
#include "td/utils/logging.h"
#include "td/utils/optional.h"
#include "td/utils/port/detail/PollableFd.h"
#include "td/utils/port/thread_local.h"
#include "td/utils/port/UdpSocketFd.h"
#include "td/utils/Span.h"
#include "td/utils/Status.h"
#include "td/utils/VectorQueue.h"

#include <array>

namespace td {

#if TD_PORT_POSIX
namespace detail {
class UdpWriter {
 public:
  static Status write_once(UdpSocketFd &fd, VectorQueue<UdpMessage> &queue) TD_WARN_UNUSED_RESULT {
    std::array<UdpSocketFd::OutboundMessage, 16> messages;
    auto to_send = queue.as_span();
    size_t to_send_n = td::min(messages.size(), to_send.size());
    to_send.truncate(to_send_n);
    for (size_t i = 0; i < to_send_n; i++) {
      messages[i].to = &to_send[i].address;
      messages[i].data = to_send[i].data.as_slice();
    }

    size_t cnt;
    auto status = fd.send_messages(Span<UdpSocketFd::OutboundMessage>(messages).truncate(to_send_n), cnt);
    queue.pop_n(cnt);
    return status;
  }
};

class UdpReaderHelper {
 public:
  void init_inbound_message(UdpSocketFd::InboundMessage &message) {
    message.from = &message_.address;
    message.error = &message_.error;
    if (buffer_.size() < MAX_PACKET_SIZE) {
      buffer_ = BufferSlice(RESERVED_SIZE);
    }
    CHECK(buffer_.size() >= MAX_PACKET_SIZE);
    message.data = buffer_.as_mutable_slice().substr(0, MAX_PACKET_SIZE);
  }

  UdpMessage extract_udp_message(UdpSocketFd::InboundMessage &message) {
    message_.data = buffer_.from_slice(message.data);
    auto size = message_.data.size();
    size = (size + 7) & ~7;
    CHECK(size <= MAX_PACKET_SIZE);
    buffer_.confirm_read(size);
    return std::move(message_);
  }

 private:
  static constexpr size_t MAX_PACKET_SIZE = 2048;
  static constexpr size_t RESERVED_SIZE = MAX_PACKET_SIZE * 8;
  UdpMessage message_;
  BufferSlice buffer_;
};

// One for thread is enough
class UdpReader {
 public:
  UdpReader() {
    for (size_t i = 0; i < messages_.size(); i++) {
      helpers_[i].init_inbound_message(messages_[i]);
    }
  }
  Status read_once(UdpSocketFd &fd, VectorQueue<UdpMessage> &queue) TD_WARN_UNUSED_RESULT {
    for (auto &message : messages_) {
      CHECK(message.data.size() == 2048);
    }
    size_t cnt = 0;
    auto status = fd.receive_messages(messages_, cnt);
    for (size_t i = 0; i < cnt; i++) {
      queue.push(helpers_[i].extract_udp_message(messages_[i]));
      helpers_[i].init_inbound_message(messages_[i]);
    }
    for (size_t i = cnt; i < messages_.size(); i++) {
      LOG_CHECK(messages_[i].data.size() == 2048)
          << " cnt = " << cnt << " i = " << i << " size = " << messages_[i].data.size() << " status = " << status;
    }
    if (status.is_error() && !UdpSocketFd::is_critical_read_error(status)) {
      queue.push(UdpMessage{{}, {}, std::move(status)});
      status = Status::OK();
    }
    return status;
  }

 private:
  static constexpr size_t BUFFER_SIZE = 16;
  std::array<UdpSocketFd::InboundMessage, BUFFER_SIZE> messages_;
  std::array<UdpReaderHelper, BUFFER_SIZE> helpers_;
};

}  // namespace detail

#endif

class BufferedUdp final : public UdpSocketFd {
 public:
  explicit BufferedUdp(UdpSocketFd fd) : UdpSocketFd(std::move(fd)) {
  }

#if TD_PORT_POSIX
  void sync_with_poll() {
    ::td::sync_with_poll(*this);
  }
  Result<optional<UdpMessage>> receive() {
    if (input_.empty() && can_read_local(*this)) {
      TRY_STATUS(flush_read_once());
    }
    if (input_.empty()) {
      return optional<UdpMessage>();
    }
    return input_.pop();
  }

  void send(UdpMessage message) {
    output_.push(std::move(message));
  }

  Status flush_send() {
    Status status;
    while (status.is_ok() && can_write_local(*this) && !output_.empty()) {
      status = flush_send_once();
    }
    return status;
  }
#endif

  UdpSocketFd move_as_udp_socket_fd() {
    return std::move(as_fd());
  }

  UdpSocketFd &as_fd() {
    return *static_cast<UdpSocketFd *>(this);
  }

 private:
#if TD_PORT_POSIX
  VectorQueue<UdpMessage> input_;
  VectorQueue<UdpMessage> output_;

  VectorQueue<UdpMessage> &input() {
    return input_;
  }
  VectorQueue<UdpMessage> &output() {
    return output_;
  }

  Status flush_send_once() TD_WARN_UNUSED_RESULT {
    return detail::UdpWriter::write_once(as_fd(), output_);
  }

  Status flush_read_once() TD_WARN_UNUSED_RESULT {
    init_thread_local<detail::UdpReader>(udp_reader_);
    return udp_reader_->read_once(as_fd(), input_);
  }

  static TD_THREAD_LOCAL detail::UdpReader *udp_reader_;
#endif
};

}  // namespace td