YugabyteDB (2.13.0.0-b42, bfc6a6643e7399ac8a0e81d06a3ee6d6571b33ab)

// Copyright (c) YugaByte, Inc.

//

// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except

// in compliance with the License.  You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software distributed under the License

// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express

// or implied.  See the License for the specific language governing permissions and limitations

// under the License.

//

#include "yb/rpc/tcp_stream.h"

#include "yb/rpc/outbound_data.h"

#include "yb/rpc/rpc_introspection.pb.h"

#include "yb/rpc/rpc_util.h"

#include "yb/util/errno.h"

#include "yb/util/flag_tags.h"

#include "yb/util/logging.h"

#include "yb/util/memory/memory_usage.h"

#include "yb/util/metrics.h"

#include "yb/util/result.h"

#include "yb/util/status_log.h"

#include "yb/util/string_util.h"

using namespace std::literals;

DECLARE_uint64(rpc_connection_timeout_ms);

DEFINE_test_flag(int32, delay_connect_ms, 0,

                 "Delay connect in tests for specified amount of milliseconds.");

METRIC_DEFINE_simple_counter(

  server, tcp_bytes_sent, "Bytes sent over TCP connections", yb::MetricUnit::kBytes);

METRIC_DEFINE_simple_counter(

  server, tcp_bytes_received, "Bytes received via TCP connections", yb::MetricUnit::kBytes);

namespace yb {

namespace rpc {

namespace {

const size_t kMaxIov = 16;

}

TcpStream::TcpStream(const StreamCreateData& data)

    : socket_(std::move(*data.socket)),

      remote_(data.remote) {

  if (data.mem_tracker) {

    mem_tracker_ = MemTracker::FindOrCreateTracker("Sending", data.mem_tracker);

  }

  if (data.metric_entity) {

    bytes_received_counter_ = METRIC_tcp_bytes_received.Instantiate(data.metric_entity);

    bytes_sent_counter_ = METRIC_tcp_bytes_sent.Instantiate(data.metric_entity);

  }

}

TcpStream::~TcpStream() {

  // Must clear the outbound_transfers_ list before deleting.

  CHECK(sending_.empty()) << ToString();

  // It's crucial that the stream is Shutdown first -- otherwise

  // our destructor will end up calling io_.stop()

  // from a possibly non-reactor thread context. This can then make all

  // hell break loose with libev.

  CHECK(!is_epoll_registered_) << ToString();

}

Status TcpStream::Start(bool connect, ev::loop_ref* loop, StreamContext* context) {

  context_ = context;

  connected_ = !connect;

  RETURN_NOT_OK(socket_.SetNoDelay(true));

  // These timeouts don't affect non-blocking sockets:

  RETURN_NOT_OK(socket_.SetSendTimeout(FLAGS_rpc_connection_timeout_ms * 1ms));

  RETURN_NOT_OK(socket_.SetRecvTimeout(FLAGS_rpc_connection_timeout_ms * 1ms));

  if (connect && FLAGS_TEST_delay_connect_ms) {

    connect_delayer_.set(*loop);

    connect_delayer_.set<TcpStream, &TcpStream::DelayConnectHandler>(this);

    connect_delayer_.start(

        static_cast<double>(FLAGS_TEST_delay_connect_ms) / MonoTime::kMillisecondsPerSecond, 0);

    return Status::OK();

  }

  return DoStart(loop, connect);

}

Status TcpStream::DoStart(ev::loop_ref* loop, bool connect) {

  if (connect) {

    auto status = socket_.Connect(remote_);

    if (!status.ok() && !status.IsTryAgain()) {

      LOG_WITH_PREFIX(WARNING) << "Connect failed: " << status;

      return status;

    }

  }

  RETURN_NOT_OK(socket_.GetSocketAddress(&local_));

  log_prefix_.clear();

  io_.set(*loop);

  io_.set<TcpStream, &TcpStream::Handler>(this);

  int events = ev::READ | (!connected_ ? ev::WRITE : 0);

  io_.start(socket_.GetFd(), events);

  DVLOG_WITH_PREFIX(3) << "Starting, listen events: " << events << ", fd: " << socket_.GetFd();

  is_epoll_registered_ = true;

  if (connected_) {

    context_->Connected();

  }

  return Status::OK();

}

void TcpStream::DelayConnectHandler(ev::timer& watcher, int revents) { // NOLINT

  if (EV_ERROR & revents) {

    LOG_WITH_PREFIX(WARNING) << "Got an error in handle delay connect";

    return;

  }

  auto status = DoStart(&watcher.loop, true /* connect */);

  if (!status.ok()) {

    Shutdown(status);

  }

}

void TcpStream::Close() {

  if (socket_.GetFd() >= 0) {

    auto status = socket_.Shutdown(true, true);

    LOG_IF(INFO, !status.ok()) << "Failed to shutdown socket: " << status;

  }

}

void TcpStream::Shutdown(const Status& status) {

  ClearSending(status);

  if (!ReadBuffer().Empty()) {

    LOG_WITH_PREFIX(WARNING) << "Shutting down with pending inbound data ("

                             << ReadBuffer().ToString() << ", status = " << status << ")";

  }

  io_.stop();

  is_epoll_registered_ = false;

  ReadBuffer().Reset();

  WARN_NOT_OK(socket_.Close(), "Error closing socket");

}

Status TcpStream::TryWrite() {

  auto result = DoWrite();

  if (result.ok()) {

    UpdateEvents();

  }

  return result;

}

TcpStream::FillIovResult TcpStream::FillIov(iovec* out) {

  int index = 0;

  size_t offset = send_position_;

  bool only_heartbeats = true;

  for (auto& data : sending_) {

    const auto wrapped_data = data.data;

    if (wrapped_data && !wrapped_data->IsHeartbeat()) {

      only_heartbeats = false;

    }

    if (data.skipped || (offset == 0 && wrapped_data && wrapped_data->IsFinished())) {

      queued_bytes_to_send_ -= data.bytes_size();

      data.ClearBytes();

      data.skipped = true;

      continue;

    }

    for (const auto& bytes : data.bytes) {

      if (offset >= bytes.size()) {

        offset -= bytes.size();

        continue;

      }

      out[index].iov_base = bytes.data() + offset;

      out[index].iov_len = bytes.size() - offset;

      offset = 0;

      if (++index == kMaxIov) {

        return FillIovResult{index, only_heartbeats};

      }

    }

  }

  return FillIovResult{index, only_heartbeats};

}

Status TcpStream::DoWrite() {

  DVLOG_WITH_PREFIX(5) << "sending_.size(): " << sending_.size();

  if (!connected_ || waiting_write_ready_ || !is_epoll_registered_) {

    DVLOG_WITH_PREFIX(5)

        << "connected_: " << connected_

        << " waiting_write_ready_: " << waiting_write_ready_

        << " is_epoll_registered_: " << is_epoll_registered_;

    return Status::OK();

  }

  // If we weren't waiting write to be ready, we could try to write data to socket.

  while (!sending_.empty()) {

    iovec iov[kMaxIov];

    auto fill_result = FillIov(iov);

    if (!fill_result.only_heartbeats) {

      context_->UpdateLastActivity();

    }

    auto result = fill_result.len != 0

        ? socket_.Writev(iov, fill_result.len)

        : 0;

    DVLOG_WITH_PREFIX(4) << "Queued writes " << queued_bytes_to_send_ << " bytes. Result "

                         << result << ", sending_.size(): " << sending_.size();

    if (PREDICT_FALSE(!result.ok())) {

      if (!result.status().IsTryAgain()) {

        YB_LOG_WITH_PREFIX_EVERY_N(WARNING, 50) << "Send failed: " << result.status();

        return result.status();

      } else {

        VLOG_WITH_PREFIX(3) << "Send temporary failed: " << result.status();

        return Status::OK();

      }

    }

    context_->UpdateLastWrite();

    IncrementCounterBy(bytes_sent_counter_, *result);

    send_position_ += *result;

    while (!sending_.empty()) {

      auto& front = sending_.front();

      size_t full_size = front.bytes_size();

      if (front.skipped) {

        PopSending();

        continue;

      }

      if (send_position_ < full_size) {

        break;

      }

      auto data = front.data;

      send_position_ -= full_size;

      PopSending();

      if (data) {

        context_->Transferred(data, Status::OK());

      }

    }

  }

  return Status::OK();

}

void TcpStream::PopSending() {

  queued_bytes_to_send_ -= sending_.front().bytes_size();

  sending_.pop_front();

  ++data_blocks_sent_;

}

void TcpStream::Handler(ev::io& watcher, int revents) {  // NOLINT

  DVLOG_WITH_PREFIX(4) << "Handler(revents=" << revents << ")";

  Status status = Status::OK();

  if (revents & ev::ERROR) {

    status = STATUS(NetworkError, ToString() + ": Handler encountered an error");

    VLOG_WITH_PREFIX(3) << status;

  }

  if (status.ok() && (revents & ev::READ)) {

    status = ReadHandler();

    if (!status.ok()) {

      VLOG_WITH_PREFIX(3) << "ReadHandler() returned error: " << status;

    }

  }

  if (status.ok() && (revents & ev::WRITE)) {

    bool just_connected = !connected_;

    if (just_connected) {

      connected_ = true;

      context_->Connected();

    }

    status = WriteHandler(just_connected);

    if (!status.ok()) {

      VLOG_WITH_PREFIX(3) << "WriteHandler() returned error: " << status;

    }

  }

  if (status.ok()) {

    UpdateEvents();

  } else {

    context_->Destroy(status);

  }

}

void TcpStream::UpdateEvents() {

  int events = 0;

  if (!read_buffer_full_) {

    events |= ev::READ;

  }

  waiting_write_ready_ = !sending_.empty() || !connected_;

  if (waiting_write_ready_) {

    events |= ev::WRITE;

  }

  if (events) {

    io_.set(events);

  }

}

Status TcpStream::ReadHandler() {

  context_->UpdateLastRead();

  for (;;) {

    auto received = Receive();

    if (PREDICT_FALSE(!received.ok())) {

      if (Errno(received.status()) == ESHUTDOWN) {

        VLOG_WITH_PREFIX(1) << "Shut down by remote end.";

      } else {

        YB_LOG_WITH_PREFIX_EVERY_N(INFO, 50) << " Recv failed: " << received;

      }

      return received.status();

    }

    // Exit the loop if we did not receive anything.

    if (!received.get()) {

      return Status::OK();

    }

    // If we were not able to process next call exit loop.

    // If status is ok, it means that we just do not have enough data to process yet.

    auto continue_receiving = TryProcessReceived();

    if (!continue_receiving.ok()) {

      return continue_receiving.status();

    }

    if (!continue_receiving.get()) {

      return Status::OK();

    }

  }

}

Result<bool> TcpStream::Receive() {

  auto iov = ReadBuffer().PrepareAppend();

  if (!iov.ok()) {

    VLOG_WITH_PREFIX(3) << "ReadBuffer().PrepareAppend() error: " << iov.status();

    if (iov.status().IsBusy()) {

      read_buffer_full_ = true;

      return false;

    }

    return iov.status();

  }

  read_buffer_full_ = false;

  if (inbound_bytes_to_skip_ > 0) {

    auto global_skip_buffer = GetGlobalSkipBuffer();

    do {

      VLOG_WITH_PREFIX(3) << "inbound_bytes_to_skip_: " << inbound_bytes_to_skip_;

      auto nread = socket_.Recv(

          global_skip_buffer.mutable_data(),

          std::min(global_skip_buffer.size(), inbound_bytes_to_skip_));

      if (!nread.ok()) {

        VLOG_WITH_PREFIX(3) << "socket_.Recv() error: " << nread.status();

        if (nread.status().IsTryAgain()) {

          return false;

        }

        return nread.status();

      }

      IncrementCounterBy(bytes_received_counter_, *nread);

      inbound_bytes_to_skip_ -= *nread;

    } while (inbound_bytes_to_skip_ > 0);

  }

  auto nread = socket_.Recvv(iov.get_ptr());

  if (!nread.ok()) {

    DVLOG_WITH_PREFIX(3) << "socket_.Recvv() error: " << nread.status();

    if (nread.status().IsTryAgain()) {

      return false;

    }

    return nread.status();

  }

  DVLOG_WITH_PREFIX(4) << "socket_.Recvv() bytes: " << *nread;

  IncrementCounterBy(bytes_received_counter_, *nread);

  ReadBuffer().DataAppended(*nread);

  return *nread != 0;

}

void TcpStream::ParseReceived() {

  auto result = TryProcessReceived();

  if (!result.ok()) {

    context_->Destroy(result.status());

    return;

  }

  if (read_buffer_full_) {

    read_buffer_full_ = false;

    UpdateEvents();

  }

}

Result<bool> TcpStream::TryProcessReceived() {

  auto& read_buffer = ReadBuffer();

  if (!read_buffer.ReadyToRead()) {

    return false;

  }

  auto result = VERIFY_RESULT(context_->ProcessReceived(ReadBufferFull(read_buffer.Full())));

  DVLOG_WITH_PREFIX(5) << "context_->ProcessReceived result: " << AsString(result);

  LOG_IF(DFATAL, inbound_bytes_to_skip_ != 0)

      << "Expected inbound_bytes_to_skip_ to be 0 instead of " << inbound_bytes_to_skip_;

  inbound_bytes_to_skip_ = result;

  return true;

}

Status TcpStream::WriteHandler(bool just_connected) {

  waiting_write_ready_ = false;

  if (sending_.empty()) {

    LOG_IF_WITH_PREFIX(WARNING, !just_connected) <<

        "Got a ready-to-write callback, but there is nothing to write.";

    return Status::OK();

  }

  return DoWrite();

}

bool TcpStream::Idle(std::string* reason_not_idle) {

  bool result = true;

  // Check if we're in the middle of receiving something.

  if (!ReadBuffer().Empty()) {

    if (reason_not_idle) {

      AppendWithSeparator("read buffer not empty", reason_not_idle);

    }

    result = false;

  }

  // Check if we still need to send something.

  if (!sending_.empty()) {

    if (reason_not_idle) {

      AppendWithSeparator("still sending", reason_not_idle);

    }

    result = false;

  }

  return result;

}

void TcpStream::ClearSending(const Status& status) {

  // Clear any outbound transfers.

  for (auto& data : sending_) {

    if (data.data) {

      context_->Transferred(data.data, status);

    }

  }

  sending_.clear();

  queued_bytes_to_send_ = 0;

}

Result<size_t> TcpStream::Send(OutboundDataPtr data) {

  // In case of TcpStream handle is absolute index of data block, since stream start.

  // So it could be cacluated as index in sending_ plus number of data blocks that were already

  // transferred.

  size_t result = data_blocks_sent_ + sending_.size();

  DVLOG_WITH_PREFIX(6) << "TcpStream::Send queueing: " << AsString(*data);

  // Serialize the actual bytes to be put on the wire.

  sending_.emplace_back(std::move(data), mem_tracker_);

  queued_bytes_to_send_ += sending_.back().bytes_size();

  DVLOG_WITH_PREFIX(4) << "Queued data, sending_.size(): " << sending_.size()

                       << ", queued_bytes_to_send_: " << queued_bytes_to_send_;

  return result;

}

void TcpStream::Cancelled(size_t handle) {

  if (handle < data_blocks_sent_) {

    return;

  }

  handle -= data_blocks_sent_;

  LOG_IF_WITH_PREFIX(DFATAL, !sending_[handle].data->IsFinished())

      << "Cancelling not finished data: " << sending_[handle].data->ToString();

  auto& entry = sending_[handle];

  if (handle == 0 && send_position_ > 0) {

    // Transfer already started, cannot drop it.

    return;

  }

  queued_bytes_to_send_ -= entry.bytes_size();

  entry.ClearBytes();

}

void TcpStream::DumpPB(const DumpRunningRpcsRequestPB& req, RpcConnectionPB* resp) {

  auto call_in_flight = resp->add_calls_in_flight();

  uint64_t sending_bytes = 0;

  for (auto& entry : sending_) {

    auto entry_bytes_size = entry.bytes_size();;

    sending_bytes += entry_bytes_size;

    if (!entry.data) {

      continue;

    }

    if (entry.data->DumpPB(req, call_in_flight)) {

      call_in_flight->set_sending_bytes(entry_bytes_size);

      call_in_flight = resp->add_calls_in_flight();

    }

  }

  resp->set_sending_bytes(sending_bytes);

  resp->mutable_calls_in_flight()->DeleteSubrange(resp->calls_in_flight_size() - 1, 1);

}

const Protocol* TcpStream::StaticProtocol() {

  static Protocol result("tcp");

  return &result;

}

StreamFactoryPtr TcpStream::Factory() {

  class TcpStreamFactory : public StreamFactory {

   private:

    std::unique_ptr<Stream> Create(const StreamCreateData& data) override {

      return std::make_unique<TcpStream>(data);

    }

  };

  return std::make_shared<TcpStreamFactory>();

}

TcpStreamSendingData::TcpStreamSendingData(OutboundDataPtr data_, const MemTrackerPtr& mem_tracker)

    : data(std::move(data_)) {

  data->Serialize(&bytes);

  if (mem_tracker) {

    size_t memory_used = sizeof(*this);

    memory_used += DynamicMemoryUsageOf(data);

    // We don't need to account `bytes` dynamic memory usage, because it stores RefCntBuffer

    // instance in internal memory and RefCntBuffer instance is referring to the same dynamic memory

    // as `data`.

    consumption = ScopedTrackedConsumption(mem_tracker, memory_used);

  }

}

} // namespace rpc

} // namespace yb