<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 <thrift/lib/cpp2/transport/rocket/server/RocketServerConnection.h>

#include <memory>
#include <utility>

#include <fmt/core.h>
#include <folly/ExceptionString.h>
#include <folly/ExceptionWrapper.h>
#include <folly/GLog.h>
#include <folly/Likely.h>
#include <folly/MapUtil.h>
#include <folly/Overload.h>
#include <folly/ScopeGuard.h>
#include <folly/SocketAddress.h>
#include <folly/Utility.h>
#include <folly/dynamic.h>
#include <folly/io/Cursor.h>
#include <folly/io/IOBuf.h>
#include <folly/io/SocketOptionMap.h>
#include <folly/io/async/AsyncTransport.h>
#include <folly/io/async/DelayedDestruction.h>

#include <wangle/acceptor/ConnectionManager.h>

#include <thrift/lib/cpp/TApplicationException.h>
#include <thrift/lib/cpp2/server/LoggingEvent.h>
#include <thrift/lib/cpp2/transport/rocket/PayloadUtils.h>
#include <thrift/lib/cpp2/transport/rocket/RocketException.h>
#include <thrift/lib/cpp2/transport/rocket/framing/Frames.h>
#include <thrift/lib/cpp2/transport/rocket/framing/Util.h>
#include <thrift/lib/cpp2/transport/rocket/server/RocketServerFrameContext.h>
#include <thrift/lib/cpp2/transport/rocket/server/RocketServerHandler.h>
#include <thrift/lib/cpp2/transport/rocket/server/RocketSinkClientCallback.h>
#include <thrift/lib/cpp2/transport/rocket/server/RocketStreamClientCallback.h>

namespace apache {
namespace thrift {
namespace rocket {

constexpr std::chrono::seconds RocketServerConnection::SocketDrainer::kTimeout;

RocketServerConnection::RocketServerConnection(
    folly::AsyncTransport::UniquePtr socket,
    std::unique_ptr<RocketServerHandler> frameHandler,
    MemoryTracker& ingressMemoryTracker,
    MemoryTracker& egressMemoryTracker,
    const Config& cfg)
    : evb_(*socket->getEventBase()),
      socket_(std::move(socket)),
      rawSocket_(
          socket_ ? socket_->getUnderlyingTransport<folly::AsyncSocket>()
                  : nullptr),
      frameHandler_(std::move(frameHandler)),
      streamStarvationTimeout_(cfg.streamStarvationTimeout),
      egressBufferBackpressureThreshold_(cfg.egressBufferBackpressureThreshold),
      egressBufferRecoverySize_(
          cfg.egressBufferBackpressureThreshold *
          cfg.egressBufferBackpressureRecoveryFactor),
      writeBatcher_(*this, cfg.writeBatchingInterval, cfg.writeBatchingSize),
      socketDrainer_(*this),
      ingressMemoryTracker_(ingressMemoryTracker),
      egressMemoryTracker_(egressMemoryTracker) {
  CHECK(socket_);
  CHECK(frameHandler_);
  socket_->setReadCB(&parser_);
  if (rawSocket_) {
    rawSocket_->setBufferCallback(this);
    rawSocket_->setSendTimeout(cfg.socketWriteTimeout.count());
  }
}

RocketStreamClientCallback& RocketServerConnection::createStreamClientCallback(
    StreamId streamId,
    RocketServerConnection& connection,
    uint32_t initialRequestN) {
  auto callback = std::make_unique<RocketStreamClientCallback>(
      streamId, connection, initialRequestN);
  auto& callbackRef = *callback;
  streams_.emplace(streamId, std::move(callback));
  return callbackRef;
}

RocketSinkClientCallback& RocketServerConnection::createSinkClientCallback(
    StreamId streamId, RocketServerConnection& connection) {
  auto callback =
      std::make_unique<RocketSinkClientCallback>(streamId, connection);
  auto& callbackRef = *callback;
  streams_.emplace(streamId, std::move(callback));
  return callbackRef;
}

void RocketServerConnection::flushWrites(
    std::unique_ptr<folly::IOBuf> writes, WriteBatchContext&& context) {
  DestructorGuard dg(this);
  DVLOG(10) << fmt::format("write: {} B", writes->computeChainDataLength());
  inflightWritesQueue_.push_back(std::move(context));
  socket_->writeChain(this, std::move(writes));
}

void RocketServerConnection::send(
    std::unique_ptr<folly::IOBuf> data,
    apache::thrift::MessageChannel::SendCallbackPtr cb) {
  evb_.dcheckIsInEventBaseThread();

  if (state_ != ConnectionState::ALIVE && state_ != ConnectionState::DRAINING) {
    return;
  }

  writeBatcher_.enqueueWrite(std::move(data), std::move(cb));
}

RocketServerConnection::~RocketServerConnection() {
  DCHECK(inflightRequests_ == 0);
  DCHECK(inflightWritesQueue_.empty());
  DCHECK(inflightSinkFinalResponses_ == 0);
  DCHECK(writeBatcher_.empty());
  DCHECK(activePausedHandlers_ == 0);
  if (rawSocket_) {
    rawSocket_->setBufferCallback(nullptr);
  }
  socket_.reset();
  // reclaim any memory in use by pending writes
  if (egressBufferSize_) {<commit_msg>Do not manually reset `RocketServerConnection::socket_`<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 <thrift/lib/cpp2/transport/rocket/server/RocketServerConnection.h>

#include <memory>
#include <utility>

#include <fmt/core.h>
#include <folly/ExceptionString.h>
#include <folly/ExceptionWrapper.h>
#include <folly/GLog.h>
#include <folly/Likely.h>
#include <folly/MapUtil.h>
#include <folly/Overload.h>
#include <folly/ScopeGuard.h>
#include <folly/SocketAddress.h>
#include <folly/Utility.h>
#include <folly/dynamic.h>
#include <folly/io/Cursor.h>
#include <folly/io/IOBuf.h>
#include <folly/io/SocketOptionMap.h>
#include <folly/io/async/AsyncTransport.h>
#include <folly/io/async/DelayedDestruction.h>

#include <wangle/acceptor/ConnectionManager.h>

#include <thrift/lib/cpp/TApplicationException.h>
#include <thrift/lib/cpp2/server/LoggingEvent.h>
#include <thrift/lib/cpp2/transport/rocket/PayloadUtils.h>
#include <thrift/lib/cpp2/transport/rocket/RocketException.h>
#include <thrift/lib/cpp2/transport/rocket/framing/Frames.h>
#include <thrift/lib/cpp2/transport/rocket/framing/Util.h>
#include <thrift/lib/cpp2/transport/rocket/server/RocketServerFrameContext.h>
#include <thrift/lib/cpp2/transport/rocket/server/RocketServerHandler.h>
#include <thrift/lib/cpp2/transport/rocket/server/RocketSinkClientCallback.h>
#include <thrift/lib/cpp2/transport/rocket/server/RocketStreamClientCallback.h>

namespace apache {
namespace thrift {
namespace rocket {

constexpr std::chrono::seconds RocketServerConnection::SocketDrainer::kTimeout;

RocketServerConnection::RocketServerConnection(
    folly::AsyncTransport::UniquePtr socket,
    std::unique_ptr<RocketServerHandler> frameHandler,
    MemoryTracker& ingressMemoryTracker,
    MemoryTracker& egressMemoryTracker,
    const Config& cfg)
    : evb_(*socket->getEventBase()),
      socket_(std::move(socket)),
      rawSocket_(
          socket_ ? socket_->getUnderlyingTransport<folly::AsyncSocket>()
                  : nullptr),
      frameHandler_(std::move(frameHandler)),
      streamStarvationTimeout_(cfg.streamStarvationTimeout),
      egressBufferBackpressureThreshold_(cfg.egressBufferBackpressureThreshold),
      egressBufferRecoverySize_(
          cfg.egressBufferBackpressureThreshold *
          cfg.egressBufferBackpressureRecoveryFactor),
      writeBatcher_(*this, cfg.writeBatchingInterval, cfg.writeBatchingSize),
      socketDrainer_(*this),
      ingressMemoryTracker_(ingressMemoryTracker),
      egressMemoryTracker_(egressMemoryTracker) {
  CHECK(socket_);
  CHECK(frameHandler_);
  socket_->setReadCB(&parser_);
  if (rawSocket_) {
    rawSocket_->setBufferCallback(this);
    rawSocket_->setSendTimeout(cfg.socketWriteTimeout.count());
  }
}

RocketStreamClientCallback& RocketServerConnection::createStreamClientCallback(
    StreamId streamId,
    RocketServerConnection& connection,
    uint32_t initialRequestN) {
  auto callback = std::make_unique<RocketStreamClientCallback>(
      streamId, connection, initialRequestN);
  auto& callbackRef = *callback;
  streams_.emplace(streamId, std::move(callback));
  return callbackRef;
}

RocketSinkClientCallback& RocketServerConnection::createSinkClientCallback(
    StreamId streamId, RocketServerConnection& connection) {
  auto callback =
      std::make_unique<RocketSinkClientCallback>(streamId, connection);
  auto& callbackRef = *callback;
  streams_.emplace(streamId, std::move(callback));
  return callbackRef;
}

void RocketServerConnection::flushWrites(
    std::unique_ptr<folly::IOBuf> writes, WriteBatchContext&& context) {
  DestructorGuard dg(this);
  DVLOG(10) << fmt::format("write: {} B", writes->computeChainDataLength());
  inflightWritesQueue_.push_back(std::move(context));
  socket_->writeChain(this, std::move(writes));
}

void RocketServerConnection::send(
    std::unique_ptr<folly::IOBuf> data,
    apache::thrift::MessageChannel::SendCallbackPtr cb) {
  evb_.dcheckIsInEventBaseThread();

  if (state_ != ConnectionState::ALIVE && state_ != ConnectionState::DRAINING) {
    return;
  }

  writeBatcher_.enqueueWrite(std::move(data), std::move(cb));
}

RocketServerConnection::~RocketServerConnection() {
  DCHECK(inflightRequests_ == 0);
  DCHECK(inflightWritesQueue_.empty());
  DCHECK(inflightSinkFinalResponses_ == 0);
  DCHECK(writeBatcher_.empty());
  DCHECK(activePausedHandlers_ == 0);
  if (rawSocket_) {
    rawSocket_->setBufferCallback(nullptr);
  }

  // Subtle: Close the socket, which will fail all outstanding writes and
  // unsubscribe the read callback, but do not destroy the object itself, since
  // other member variables of RocketServerConnection may be borrowing the
  // object.
  socket_->closeNow();

  // reclaim any memory in use by pending writes
  if (egressBufferSize_) {
{'fraction_non_alphanumeric': 0.11745545123586894, 'fraction_numerical': 0.004023759340869898, 'mean_word_length': 5.512164691203992, 'pattern_counts': {'":': 0, '<': 79, '<?xml version=': 0, '>': 92, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4824030', 'n_tokens_mistral': 3184, 'n_tokens_neox': 3198, 'n_words': 603}