+#include "remote_connection.h"

+

+#include "key_value_printer.h"

+#include "mb_lwtrace.h"

+#include "network.h"

+#include "remote_client_connection.h"

+#include "remote_client_session.h"

+#include "remote_server_session.h"

+#include "session_impl.h"

+

+#include <library/cpp/messagebus/actor/what_thread_does.h>

+

+#include <util/generic/cast.h>

+#include <util/network/init.h>

+#include <util/system/atomic.h>

+

+LWTRACE_USING(LWTRACE_MESSAGEBUS_PROVIDER)

+

+using namespace NActor;

+using namespace NBus;

+using namespace NBus::NPrivate;

+

+namespace NBus {

+ namespace NPrivate {

+ TRemoteConnection::TRemoteConnection(TRemoteSessionPtr session, ui64 connectionId, TNetAddr addr)

+ : TActor<TRemoteConnection, TWriterTag>(session->Queue->WorkQueue.Get())

+ , TActor<TRemoteConnection, TReaderTag>(session->Queue->WorkQueue.Get())

+ , TScheduleActor<TRemoteConnection, TWriterTag>(&session->Queue->Scheduler)

+ , Session(session)

+ , Proto(session->Proto)

+ , Config(session->Config)

+ , RemovedFromSession(false)

+ , ConnectionId(connectionId)

+ , PeerAddr(addr)

+ , PeerAddrSocketAddr(addr)

+ , CreatedTime(TInstant::Now())

+ , ReturnConnectFailedImmediately(false)

+ , GranStatus(Config.Secret.StatusFlushPeriod)

+ , QuotaMsg(!Session->IsSource_, Config.PerConnectionMaxInFlight, 0)

+ , QuotaBytes(!Session->IsSource_, Config.PerConnectionMaxInFlightBySize, 0)

+ , MaxBufferSize(session->Config.MaxBufferSize)

+ , ShutdownReason(MESSAGE_OK)

+ {

+ WriterData.Status.ConnectionId = connectionId;

+ WriterData.Status.PeerAddr = PeerAddr;

+ ReaderData.Status.ConnectionId = connectionId;

+

+ const TInstant now = TInstant::Now();

+

+ WriterFillStatus();

+

+ GranStatus.Writer.Update(WriterData.Status, now, true);

+ GranStatus.Reader.Update(ReaderData.Status, now, true);

+ }

+

+ TRemoteConnection::~TRemoteConnection() {

+ Y_VERIFY(ReplyQueue.IsEmpty());

+ }

+

+ TRemoteConnection::TWriterData::TWriterData()

+ : Down(0)

+ , SocketVersion(0)

+ , InFlight(0)

+ , AwakeFlags(0)

+ , State(WRITER_FILLING)

+ {

+ }

+

+ TRemoteConnection::TWriterData::~TWriterData() {

+ Y_VERIFY(AtomicGet(Down));

+ Y_VERIFY(SendQueue.Empty());

+ }

+

+ bool TRemoteConnection::TReaderData::HasBytesInBuf(size_t bytes) noexcept {

+ size_t left = Buffer.Size() - Offset;

+

+ return (MoreBytes = left >= bytes ? 0 : bytes - left) == 0;

+ }

+

+ void TRemoteConnection::TWriterData::SetChannel(NEventLoop::TChannelPtr channel) {

+ Y_VERIFY(!Channel, "must not have channel");

+ Y_VERIFY(Buffer.GetBuffer().Empty() && Buffer.LeftSize() == 0, "buffer must be empty");

+ Y_VERIFY(State == WRITER_FILLING, "state must be initial");

+ Channel = channel;

+ }

+

+ void TRemoteConnection::TReaderData::SetChannel(NEventLoop::TChannelPtr channel) {

+ Y_VERIFY(!Channel, "must not have channel");

+ Y_VERIFY(Buffer.Empty(), "buffer must be empty");

+ Channel = channel;

+ }

+

+ void TRemoteConnection::TWriterData::DropChannel() {

+ if (!!Channel) {

+ Channel->Unregister();

+ Channel.Drop();

+ }

+

+ Buffer.Reset();

+ State = WRITER_FILLING;

+ }

+

+ void TRemoteConnection::TReaderData::DropChannel() {

+ // TODO: make Drop call Unregister

+ if (!!Channel) {

+ Channel->Unregister();

+ Channel.Drop();

+ }

+ Buffer.Reset();

+ Offset = 0;

+ }

+

+ TRemoteConnection::TReaderData::TReaderData()

+ : Down(0)

+ , SocketVersion(0)

+ , Offset(0)

+ , MoreBytes(0)

+ {

+ }

+

+ TRemoteConnection::TReaderData::~TReaderData() {

+ Y_VERIFY(AtomicGet(Down));

+ }

+

+ void TRemoteConnection::Send(TNonDestroyingAutoPtr<TBusMessage> msg) {

+ BeforeSendQueue.Enqueue(msg.Release());

+ AtomicIncrement(WriterData.InFlight);

+ ScheduleWrite();

+ }

+

+ void TRemoteConnection::ClearOutgoingQueue(TMessagesPtrs& result, bool reconnect) {

+ if (!reconnect) {

+ // Do not clear send queue if reconnecting

+ WriterData.SendQueue.Clear(&result);

+ }

+ }

+

+ void TRemoteConnection::Shutdown(EMessageStatus status) {

+ ScheduleShutdown(status);

+

+ ReaderData.ShutdownComplete.WaitI();

+ WriterData.ShutdownComplete.WaitI();

+ }

+

+ void TRemoteConnection::TryConnect() {

+ Y_FAIL("TryConnect is client connection only operation");

+ }

+

+ void TRemoteConnection::ScheduleRead() {

+ GetReaderActor()->Schedule();

+ }

+

+ void TRemoteConnection::ScheduleWrite() {

+ GetWriterActor()->Schedule();

+ }

+

+ void TRemoteConnection::WriterRotateCounters() {

+ if (!WriterData.TimeToRotateCounters.FetchTask()) {

+ return;

+ }

+

+ WriterData.Status.DurationCounterPrev = WriterData.Status.DurationCounter;

+ Reset(WriterData.Status.DurationCounter);

+ }

+

+ void TRemoteConnection::WriterSendStatus(TInstant now, bool force) {

+ GranStatus.Writer.Update(std::bind(&TRemoteConnection::WriterGetStatus, this), now, force);

+ }

+

+ void TRemoteConnection::ReaderSendStatus(TInstant now, bool force) {

+ GranStatus.Reader.Update(std::bind(&TRemoteConnection::ReaderFillStatus, this), now, force);

+ }

+

+ const TRemoteConnectionReaderStatus& TRemoteConnection::ReaderFillStatus() {

+ ReaderData.Status.BufferSize = ReaderData.Buffer.Capacity();

+ ReaderData.Status.QuotaMsg = QuotaMsg.Tokens();

+ ReaderData.Status.QuotaBytes = QuotaBytes.Tokens();

+

+ return ReaderData.Status;

+ }

+

+ void TRemoteConnection::ProcessItem(TReaderTag, ::NActor::TDefaultTag, TWriterToReaderSocketMessage readSocket) {

+ if (AtomicGet(ReaderData.Down)) {

+ ReaderData.Status.Fd = INVALID_SOCKET;

+ return;

+ }

+

+ ReaderData.DropChannel();

+

+ ReaderData.Status.Fd = readSocket.Socket;

+ ReaderData.SocketVersion = readSocket.SocketVersion;

+

+ if (readSocket.Socket != INVALID_SOCKET) {

+ ReaderData.SetChannel(Session->ReadEventLoop.Register(readSocket.Socket, this, ReadCookie));

+ ReaderData.Channel->EnableRead();

+ }

+ }

+

+ void TRemoteConnection::ProcessItem(TWriterTag, TReconnectTag, ui32 socketVersion) {

+ Y_VERIFY(socketVersion <= WriterData.SocketVersion, "something weird");

+

+ if (WriterData.SocketVersion != socketVersion) {

+ return;

+ }

+ Y_VERIFY(WriterData.Status.Connected, "must be connected at this point");

+ Y_VERIFY(!!WriterData.Channel, "must have channel at this point");

+

+ WriterData.Status.Connected = false;

+ WriterData.DropChannel();

+ WriterData.Status.MyAddr = TNetAddr();

+ ++WriterData.SocketVersion;

+ LastConnectAttempt = TInstant();

+

+ TMessagesPtrs cleared;

+ ClearOutgoingQueue(cleared, true);

+ WriterErrorMessages(cleared, MESSAGE_DELIVERY_FAILED);

+

+ FireClientConnectionEvent(TClientConnectionEvent::DISCONNECTED);

+

+ ReaderGetSocketQueue()->EnqueueAndSchedule(TWriterToReaderSocketMessage(INVALID_SOCKET, WriterData.SocketVersion));

+ }

+

+ void TRemoteConnection::ProcessItem(TWriterTag, TWakeReaderTag, ui32 awakeFlags) {

+ WriterData.AwakeFlags |= awakeFlags;

+

+ ReadQuotaWakeup();

+ }

+

+ void TRemoteConnection::Act(TReaderTag) {

+ TInstant now = TInstant::Now();

+

+ ReaderData.Status.Acts += 1;

+

+ ReaderGetSocketQueue()->DequeueAllLikelyEmpty();

+

+ if (AtomicGet(ReaderData.Down)) {

+ ReaderData.DropChannel();

+

+ ReaderProcessStatusDown();

+ ReaderData.ShutdownComplete.Signal();

+

+ } else if (!!ReaderData.Channel) {

+ Y_ASSERT(ReaderData.ReadMessages.empty());

+

+ for (int i = 0;; ++i) {

+ if (i == 100) {

+ // perform other tasks

+ GetReaderActor()->AddTaskFromActorLoop();

+ break;

+ }

+

+ if (NeedInterruptRead()) {

+ ReaderData.Channel->EnableRead();

+ break;

+ }

+

+ if (!ReaderFillBuffer())

+ break;

+

+ if (!ReaderProcessBuffer())

+ break;

+ }

+

+ ReaderFlushMessages();

+ }

+

+ ReaderSendStatus(now);

+ }

+

+ bool TRemoteConnection::QuotaAcquire(size_t msg, size_t bytes) {

+ ui32 wakeFlags = 0;

+

+ if (!QuotaMsg.Acquire(msg))

+ wakeFlags |= WAKE_QUOTA_MSG;

+

+ else if (!QuotaBytes.Acquire(bytes))

+ wakeFlags |= WAKE_QUOTA_BYTES;

+

+ if (wakeFlags) {

+ ReaderData.Status.QuotaExhausted++;

+

+ WriterGetWakeQueue()->EnqueueAndSchedule(wakeFlags);

+ }

+

+ return wakeFlags == 0;

+ }

+

+ void TRemoteConnection::QuotaConsume(size_t msg, size_t bytes) {

+ QuotaMsg.Consume(msg);

+ QuotaBytes.Consume(bytes);

+ }

+

+ void TRemoteConnection::QuotaReturnSelf(size_t items, size_t bytes) {

+ if (QuotaReturnValues(items, bytes))

+ ReadQuotaWakeup();

+ }

+

+ void TRemoteConnection::QuotaReturnAside(size_t items, size_t bytes) {

+ if (QuotaReturnValues(items, bytes) && !AtomicGet(WriterData.Down))

+ WriterGetWakeQueue()->EnqueueAndSchedule(0x0);

+ }

+

+ bool TRemoteConnection::QuotaReturnValues(size_t items, size_t bytes) {

+ bool rMsg = QuotaMsg.Return(items);

+ bool rBytes = QuotaBytes.Return(bytes);

+

+ return rMsg || rBytes;

+ }

+

+ void TRemoteConnection::ReadQuotaWakeup() {

+ const ui32 mask = WriterData.AwakeFlags & WriteWakeFlags();

+

+ if (mask && mask == WriterData.AwakeFlags) {

+ WriterData.Status.ReaderWakeups++;

+ WriterData.AwakeFlags = 0;

+

+ ScheduleRead();

+ }

+ }

+

+ ui32 TRemoteConnection::WriteWakeFlags() const {

+ ui32 awakeFlags = 0;

+

+ if (QuotaMsg.IsAboveWake())

+ awakeFlags |= WAKE_QUOTA_MSG;

+

+ if (QuotaBytes.IsAboveWake())

+ awakeFlags |= WAKE_QUOTA_BYTES;

+

+ return awakeFlags;

+ }

+

+ bool TRemoteConnection::ReaderProcessBuffer() {

+ TInstant now = TInstant::Now();

+

+ for (;;) {

+ if (!ReaderData.HasBytesInBuf(sizeof(TBusHeader))) {

+ break;

+ }

+

+ TBusHeader header(MakeArrayRef(ReaderData.Buffer.Data() + ReaderData.Offset, ReaderData.Buffer.Size() - ReaderData.Offset));

+

+ if (header.Size < sizeof(TBusHeader)) {

+ LWPROBE(Error, ToString(MESSAGE_HEADER_CORRUPTED), ToString(PeerAddr), ToString(header.Size));

+ ReaderData.Status.Incremental.StatusCounter[MESSAGE_HEADER_CORRUPTED] += 1;

+ ScheduleShutdownOnServerOrReconnectOnClient(MESSAGE_HEADER_CORRUPTED, false);

+ return false;

+ }

+

+ if (!IsVersionNegotiation(header) && !IsBusKeyValid(header.Id)) {

+ LWPROBE(Error, ToString(MESSAGE_HEADER_CORRUPTED), ToString(PeerAddr), ToString(header.Size));

+ ReaderData.Status.Incremental.StatusCounter[MESSAGE_HEADER_CORRUPTED] += 1;

+ ScheduleShutdownOnServerOrReconnectOnClient(MESSAGE_HEADER_CORRUPTED, false);

+ return false;

+ }

+

+ if (header.Size > Config.MaxMessageSize) {

+ LWPROBE(Error, ToString(MESSAGE_MESSAGE_TOO_LARGE), ToString(PeerAddr), ToString(header.Size));

+ ReaderData.Status.Incremental.StatusCounter[MESSAGE_MESSAGE_TOO_LARGE] += 1;

+ ScheduleShutdownOnServerOrReconnectOnClient(MESSAGE_MESSAGE_TOO_LARGE, false);

+ return false;

+ }

+

+ if (!ReaderData.HasBytesInBuf(header.Size)) {

+ if (ReaderData.Offset == 0) {

+ ReaderData.Buffer.Reserve(header.Size);

+ }

+ break;

+ }

+

+ if (!QuotaAcquire(1, header.Size))

+ return false;

+

+ if (!MessageRead(MakeArrayRef(ReaderData.Buffer.Data() + ReaderData.Offset, header.Size), now)) {

+ return false;

+ }

+

+ ReaderData.Offset += header.Size;

+ }

+

+ ReaderData.Buffer.ChopHead(ReaderData.Offset);

+ ReaderData.Offset = 0;

+

+ if (ReaderData.Buffer.Capacity() > MaxBufferSize && ReaderData.Buffer.Size() <= MaxBufferSize) {

+ ReaderData.Status.Incremental.BufferDrops += 1;

+

+ TBuffer temp;

+ // probably should use another constant

+ temp.Reserve(Config.DefaultBufferSize);

+ temp.Append(ReaderData.Buffer.Data(), ReaderData.Buffer.Size());

+

+ ReaderData.Buffer.Swap(temp);

+ }

+

+ return true;

+ }

+

+ bool TRemoteConnection::ReaderFillBuffer() {

+ if (!ReaderData.BufferMore())

+ return true;

+

+ if (ReaderData.Buffer.Avail() == 0) {

+ if (ReaderData.Buffer.Size() == 0) {

+ ReaderData.Buffer.Reserve(Config.DefaultBufferSize);

+ } else {

+ ReaderData.Buffer.Reserve(ReaderData.Buffer.Size() * 2);

+ }

+ }

+

+ Y_ASSERT(ReaderData.Buffer.Avail() > 0);

+

+ ssize_t bytes;

+ {

+ TWhatThreadDoesPushPop pp("recv syscall");

+ bytes = SocketRecv(ReaderData.Channel->GetSocket(), TArrayRef<char>(ReaderData.Buffer.Pos(), ReaderData.Buffer.Avail()));

+ }

+

+ if (bytes < 0) {

+ if (WouldBlock()) {

+ ReaderData.Channel->EnableRead();

+ return false;

+ } else {

+ ReaderData.Channel->DisableRead();

+ ScheduleShutdownOnServerOrReconnectOnClient(MESSAGE_DELIVERY_FAILED, false);

+ return false;

+ }

+ }

+

+ if (bytes == 0) {

+ ReaderData.Channel->DisableRead();

+ // TODO: incorrect: it is possible that only input is shutdown, and output is available

+ ScheduleShutdownOnServerOrReconnectOnClient(MESSAGE_DELIVERY_FAILED, false);

+ return false;

+ }

+

+ ReaderData.Status.Incremental.NetworkOps += 1;

+

+ ReaderData.Buffer.Advance(bytes);

+ ReaderData.MoreBytes = 0;

+ return true;

+ }

+

+ void TRemoteConnection::ClearBeforeSendQueue(EMessageStatus reason) {

+ BeforeSendQueue.DequeueAll(std::bind(&TRemoteConnection::WriterBeforeWriteErrorMessage, this, std::placeholders::_1, reason));

+ }

+

+ void TRemoteConnection::ClearReplyQueue(EMessageStatus reason) {

+ TVectorSwaps<TBusMessagePtrAndHeader> replyQueueTemp;

+ Y_ASSERT(replyQueueTemp.empty());

+ ReplyQueue.DequeueAllSingleConsumer(&replyQueueTemp);

+

+ TVector<TBusMessage*> messages;

+ for (TVectorSwaps<TBusMessagePtrAndHeader>::reverse_iterator message = replyQueueTemp.rbegin();

+ message != replyQueueTemp.rend(); ++message) {

+ messages.push_back(message->MessagePtr.Release());

+ }

+

+ WriterErrorMessages(messages, reason);

+

+ replyQueueTemp.clear();

+ }

+

+ void TRemoteConnection::ProcessBeforeSendQueueMessage(TBusMessage* message, TInstant now) {

+ // legacy clients expect this field to be set

+ if (!Session->IsSource_) {

+ message->SendTime = now.MilliSeconds();

+ }

+

+ WriterData.SendQueue.PushBack(message);

+ }

+

+ void TRemoteConnection::ProcessBeforeSendQueue(TInstant now) {

+ BeforeSendQueue.DequeueAll(std::bind(&TRemoteConnection::ProcessBeforeSendQueueMessage, this, std::placeholders::_1, now));

+ }

+

+ void TRemoteConnection::WriterFillInFlight() {

+ // this is hack for TLoadBalancedProtocol

+ WriterFillStatus();

+ AtomicSet(WriterData.InFlight, WriterData.Status.GetInFlight());

+ }

+

+ const TRemoteConnectionWriterStatus& TRemoteConnection::WriterGetStatus() {

+ WriterRotateCounters();

+ WriterFillStatus();

+

+ return WriterData.Status;

+ }

+

+ void TRemoteConnection::WriterFillStatus() {

+ if (!!WriterData.Channel) {

+ WriterData.Status.Fd = WriterData.Channel->GetSocket();

+ } else {

+ WriterData.Status.Fd = INVALID_SOCKET;

+ }

+ WriterData.Status.BufferSize = WriterData.Buffer.Capacity();

+ WriterData.Status.SendQueueSize = WriterData.SendQueue.Size();

+ WriterData.Status.State = WriterData.State;

+ }

+

+ void TRemoteConnection::WriterProcessStatusDown() {

+ Session->GetDeadConnectionWriterStatusQueue()->EnqueueAndSchedule(WriterData.Status.Incremental);

+ Reset(WriterData.Status.Incremental);

+ }

+

+ void TRemoteConnection::ReaderProcessStatusDown() {

+ Session->GetDeadConnectionReaderStatusQueue()->EnqueueAndSchedule(ReaderData.Status.Incremental);

+ Reset(ReaderData.Status.Incremental);

+ }

+

+ void TRemoteConnection::ProcessWriterDown() {

+ if (!RemovedFromSession) {

+ Session->GetRemoveConnectionQueue()->EnqueueAndSchedule(this);

+

+ if (Session->IsSource_) {

+ if (WriterData.Status.Connected) {

+ FireClientConnectionEvent(TClientConnectionEvent::DISCONNECTED);

+ }

+ }

+

+ LWPROBE(Disconnected, ToString(PeerAddr));

+ RemovedFromSession = true;

+ }

+

+ WriterData.DropChannel();

+

+ DropEnqueuedData(ShutdownReason, MESSAGE_SHUTDOWN);

+

+ WriterProcessStatusDown();

+

+ WriterData.ShutdownComplete.Signal();

+ }

+

+ void TRemoteConnection::DropEnqueuedData(EMessageStatus reason, EMessageStatus reasonForQueues) {

+ ClearReplyQueue(reasonForQueues);

+ ClearBeforeSendQueue(reasonForQueues);

+ WriterGetReconnectQueue()->Clear();

+ WriterGetWakeQueue()->Clear();

+

+ TMessagesPtrs cleared;

+ ClearOutgoingQueue(cleared, false);

+

+ if (!Session->IsSource_) {

+ for (auto& i : cleared) {

+ TBusMessagePtrAndHeader h(i);

+ CheckedCast<TRemoteServerSession*>(Session.Get())->ReleaseInWorkResponses(MakeArrayRef(&h, 1));

+ // assignment back is weird

+ i = h.MessagePtr.Release();

+ // and this part is not batch

+ }

+ }

+

+ WriterErrorMessages(cleared, reason);

+ }

+

+ void TRemoteConnection::BeforeTryWrite() {

+ }

+

+ void TRemoteConnection::Act(TWriterTag) {

+ TInstant now = TInstant::Now();

+

+ WriterData.Status.Acts += 1;

+

+ if (Y_UNLIKELY(AtomicGet(WriterData.Down))) {

+ // dump status must work even if WriterDown

+ WriterSendStatus(now, true);

+ ProcessWriterDown();

+ return;

+ }

+

+ ProcessBeforeSendQueue(now);

+

+ BeforeTryWrite();

+

+ WriterFillInFlight();

+

+ WriterGetReconnectQueue()->DequeueAllLikelyEmpty();

+

+ if (!WriterData.Status.Connected) {

+ TryConnect();

+ } else {

+ for (int i = 0;; ++i) {

+ if (i == 100) {

+ // perform other tasks

+ GetWriterActor()->AddTaskFromActorLoop();

+ break;

+ }

+

+ if (WriterData.State == WRITER_FILLING) {

+ WriterFillBuffer();

+

+ if (WriterData.State == WRITER_FILLING) {

+ WriterData.Channel->DisableWrite();

+ break;

+ }

+

+ Y_ASSERT(!WriterData.Buffer.Empty());

+ }

+

+ if (WriterData.State == WRITER_FLUSHING) {

+ WriterFlushBuffer();

+

+ if (WriterData.State == WRITER_FLUSHING) {

+ break;

+ }

+ }

+ }

+ }

+

+ WriterGetWakeQueue()->DequeueAllLikelyEmpty();

+

+ WriterSendStatus(now);

+ }

+

+ void TRemoteConnection::WriterFlushBuffer() {

+ Y_ASSERT(WriterData.State == WRITER_FLUSHING);

+ Y_ASSERT(!WriterData.Buffer.Empty());

+

+ WriterData.CorkUntil = TInstant::Zero();

+

+ while (!WriterData.Buffer.Empty()) {

+ ssize_t bytes;

+ {

+ TWhatThreadDoesPushPop pp("send syscall");

+ bytes = SocketSend(WriterData.Channel->GetSocket(), TArrayRef<const char>(WriterData.Buffer.LeftPos(), WriterData.Buffer.Size()));

+ }

+

+ if (bytes < 0) {

+ if (WouldBlock()) {

+ WriterData.Channel->EnableWrite();

+ return;

+ } else {

+ WriterData.Channel->DisableWrite();

+ ScheduleShutdownOnServerOrReconnectOnClient(MESSAGE_DELIVERY_FAILED, true);

+ return;

+ }

+ }

+

+ WriterData.Status.Incremental.NetworkOps += 1;

+

+ WriterData.Buffer.LeftProceed(bytes);

+ }

+

+ WriterData.Buffer.Clear();

+ if (WriterData.Buffer.Capacity() > MaxBufferSize) {

+ WriterData.Status.Incremental.BufferDrops += 1;

+ WriterData.Buffer.Reset();

+ }

+

+ WriterData.State = WRITER_FILLING;

+ }

+

+ void TRemoteConnection::ScheduleShutdownOnServerOrReconnectOnClient(EMessageStatus status, bool writer) {

+ if (Session->IsSource_) {

+ WriterGetReconnectQueue()->EnqueueAndSchedule(writer ? WriterData.SocketVersion : ReaderData.SocketVersion);

+ } else {

+ ScheduleShutdown(status);

+ }

+ }

+

+ void TRemoteConnection::ScheduleShutdown(EMessageStatus status) {

+ ShutdownReason = status;

+

+ AtomicSet(ReaderData.Down, 1);

+ ScheduleRead();

+

+ AtomicSet(WriterData.Down, 1);

+ ScheduleWrite();

+ }

+

+ void TRemoteConnection::CallSerialize(TBusMessage* msg, TBuffer& buffer) const {

+ size_t posForAssertion = buffer.Size();

+ Proto->Serialize(msg, buffer);

+ Y_VERIFY(buffer.Size() >= posForAssertion,

+ "incorrect Serialize implementation, pos before serialize: %d, pos after serialize: %d",

+ int(posForAssertion), int(buffer.Size()));

+ }

+

+ namespace {

+ inline void WriteHeader(const TBusHeader& header, TBuffer& data) {

+ data.Reserve(data.Size() + sizeof(TBusHeader));

+ /// \todo hton instead of memcpy

+ memcpy(data.Data() + data.Size(), &header, sizeof(TBusHeader));

+ data.Advance(sizeof(TBusHeader));

+ }

+

+ inline void WriteDummyHeader(TBuffer& data) {

+ data.Resize(data.Size() + sizeof(TBusHeader));

+ }

+

+ }

+

+ void TRemoteConnection::SerializeMessage(TBusMessage* msg, TBuffer* data, TMessageCounter* counter) const {

+ size_t pos = data->Size();

+

+ size_t dataSize;

+

+ bool compressionRequested = msg->IsCompressed();

+

+ if (compressionRequested) {

+ TBuffer compdata;

+ TBuffer plaindata;

+ CallSerialize(msg, plaindata);

+

+ dataSize = sizeof(TBusHeader) + plaindata.Size();

+

+ NCodecs::TCodecPtr c = Proto->GetTransportCodec();

+ c->Encode(TStringBuf{plaindata.data(), plaindata.size()}, compdata);

+

+ if (compdata.Size() < plaindata.Size()) {

+ plaindata.Clear();

+ msg->GetHeader()->Size = sizeof(TBusHeader) + compdata.Size();

+ WriteHeader(*msg->GetHeader(), *data);

+ data->Append(compdata.Data(), compdata.Size());

+ } else {

+ compdata.Clear();

+ msg->SetCompressed(false);

+ msg->GetHeader()->Size = sizeof(TBusHeader) + plaindata.Size();

+ WriteHeader(*msg->GetHeader(), *data);

+ data->Append(plaindata.Data(), plaindata.Size());

+ }

+ } else {

+ WriteDummyHeader(*data);

+ CallSerialize(msg, *data);

+

+ dataSize = msg->GetHeader()->Size = data->Size() - pos;

+

+ data->Proceed(pos);

+ WriteHeader(*msg->GetHeader(), *data);

+ data->Proceed(pos + msg->GetHeader()->Size);

+ }

+

+ Y_ASSERT(msg->GetHeader()->Size == data->Size() - pos);

+ counter->AddMessage(dataSize, data->Size() - pos, msg->IsCompressed(), compressionRequested);

+ }

+

+ TBusMessage* TRemoteConnection::DeserializeMessage(TArrayRef<const char> dataRef, const TBusHeader* header, TMessageCounter* messageCounter, EMessageStatus* status) const {

+ size_t dataSize;

+

+ TBusMessage* message;

+ if (header->FlagsInternal & MESSAGE_COMPRESS_INTERNAL) {

+ TBuffer msg;

+ {

+ TBuffer plaindata;

+ NCodecs::TCodecPtr c = Proto->GetTransportCodec();

+ try {

+ TArrayRef<const char> payload = TBusMessage::GetPayload(dataRef);

+ c->Decode(TStringBuf{payload.data(), payload.size()}, plaindata);

+ } catch (...) {

+ // catch all, because

+ // http://nga.at.yandex-team.ru/replies.xml?item_no=3884

+ *status = MESSAGE_DECOMPRESS_ERROR;

+ return nullptr;

+ }

+

+ msg.Append(dataRef.data(), sizeof(TBusHeader));

+ msg.Append(plaindata.Data(), plaindata.Size());

+ }

+ TArrayRef<const char> msgRef(msg.Data(), msg.Size());

+ dataSize = sizeof(TBusHeader) + msgRef.size();

+ // TODO: track error types

+ message = Proto->Deserialize(header->Type, msgRef.Slice(sizeof(TBusHeader))).Release();

+ if (!message) {

+ *status = MESSAGE_DESERIALIZE_ERROR;

+ return nullptr;

+ }

+ *message->GetHeader() = *header;

+ message->SetCompressed(true);

+ } else {

+ dataSize = dataRef.size();

+ message = Proto->Deserialize(header->Type, dataRef.Slice(sizeof(TBusHeader))).Release();

+ if (!message) {

+ *status = MESSAGE_DESERIALIZE_ERROR;

+ return nullptr;

+ }

+ *message->GetHeader() = *header;

+ }

+

+ messageCounter->AddMessage(dataSize, dataRef.size(), header->FlagsInternal & MESSAGE_COMPRESS_INTERNAL, false);

+

+ return message;

+ }

+

+ void TRemoteConnection::ResetOneWayFlag(TArrayRef<TBusMessage*> messages) {

+ for (auto message : messages) {

+ message->LocalFlags &= ~MESSAGE_ONE_WAY_INTERNAL;

+ }

+ }

+

+ void TRemoteConnection::ReaderFlushMessages() {

+ if (!ReaderData.ReadMessages.empty()) {

+ Session->OnMessageReceived(this, ReaderData.ReadMessages);

+ ReaderData.ReadMessages.clear();

+ }

+ }

+

+ // @return false if actor should break

+ bool TRemoteConnection::MessageRead(TArrayRef<const char> readDataRef, TInstant now) {

+ TBusHeader header(readDataRef);

+

+ Y_ASSERT(readDataRef.size() == header.Size);

+

+ if (header.GetVersionInternal() != YBUS_VERSION) {

+ ReaderProcessMessageUnknownVersion(readDataRef);

+ return true;

+ }

+

+ EMessageStatus deserializeFailureStatus = MESSAGE_OK;

+ TBusMessage* r = DeserializeMessage(readDataRef, &header, &ReaderData.Status.Incremental.MessageCounter, &deserializeFailureStatus);

+

+ if (!r) {

+ Y_VERIFY(deserializeFailureStatus != MESSAGE_OK, "state check");

+ LWPROBE(Error, ToString(deserializeFailureStatus), ToString(PeerAddr), "");

+ ReaderData.Status.Incremental.StatusCounter[deserializeFailureStatus] += 1;

+ ScheduleShutdownOnServerOrReconnectOnClient(deserializeFailureStatus, false);

+ return false;

+ }

+

+ LWPROBE(Read, r->GetHeader()->Size);

+

+ r->ReplyTo = PeerAddrSocketAddr;

+

+ TBusMessagePtrAndHeader h(r);

+ r->RecvTime = now;

+

+ QuotaConsume(1, header.Size);

+

+ ReaderData.ReadMessages.push_back(h);

+ if (ReaderData.ReadMessages.size() >= 100) {

+ ReaderFlushMessages();

+ }

+

+ return true;

+ }

+

+ void TRemoteConnection::WriterFillBuffer() {

+ Y_ASSERT(WriterData.State == WRITER_FILLING);

+

+ Y_ASSERT(WriterData.Buffer.LeftSize() == 0);

+

+ if (Y_UNLIKELY(!WrongVersionRequests.IsEmpty())) {

+ TVector<TBusHeader> headers;

+ WrongVersionRequests.DequeueAllSingleConsumer(&headers);

+ for (TVector<TBusHeader>::reverse_iterator header = headers.rbegin();

+ header != headers.rend(); ++header) {

+ TBusHeader response = *header;

+ response.SendTime = NBus::Now();

+ response.Size = sizeof(TBusHeader);

+ response.FlagsInternal = 0;

+ response.SetVersionInternal(YBUS_VERSION);

+ WriteHeader(response, WriterData.Buffer.GetBuffer());

+ }

+

+ Y_ASSERT(!WriterData.Buffer.Empty());

+ WriterData.State = WRITER_FLUSHING;

+ return;

+ }

+

+ TTempTlsVector<TBusMessagePtrAndHeader, void, TVectorSwaps> writeMessages;

+

+ for (;;) {

+ THolder<TBusMessage> writeMessage(WriterData.SendQueue.PopFront());

+ if (!writeMessage) {

+ break;

+ }

+

+ if (Config.Cork != TDuration::Zero()) {

+ if (WriterData.CorkUntil == TInstant::Zero()) {

+ WriterData.CorkUntil = TInstant::Now() + Config.Cork;

+ }

+ }

+

+ size_t sizeBeforeSerialize = WriterData.Buffer.Size();

+

+ TMessageCounter messageCounter = WriterData.Status.Incremental.MessageCounter;

+

+ SerializeMessage(writeMessage.Get(), &WriterData.Buffer.GetBuffer(), &messageCounter);

+

+ size_t written = WriterData.Buffer.Size() - sizeBeforeSerialize;

+ if (written > Config.MaxMessageSize) {

+ WriterData.Buffer.GetBuffer().EraseBack(written);

+ WriterBeforeWriteErrorMessage(writeMessage.Release(), MESSAGE_MESSAGE_TOO_LARGE);

+ continue;

+ }

+

+ WriterData.Status.Incremental.MessageCounter = messageCounter;

+

+ TBusMessagePtrAndHeader h(writeMessage.Release());

+ writeMessages.GetVector()->push_back(h);

+

+ Y_ASSERT(!WriterData.Buffer.Empty());

+ if (WriterData.Buffer.Size() >= Config.SendThreshold) {

+ break;

+ }

+ }

+

+ if (!WriterData.Buffer.Empty()) {

+ if (WriterData.Buffer.Size() >= Config.SendThreshold) {

+ WriterData.State = WRITER_FLUSHING;

+ } else if (WriterData.CorkUntil == TInstant::Zero()) {

+ WriterData.State = WRITER_FLUSHING;

+ } else if (TInstant::Now() >= WriterData.CorkUntil) {

+ WriterData.State = WRITER_FLUSHING;

+ } else {

+ // keep filling

+ Y_ASSERT(WriterData.State == WRITER_FILLING);

+ GetWriterSchedulerActor()->ScheduleAt(WriterData.CorkUntil);

+ }

+ } else {

+ // keep filling

+ Y_ASSERT(WriterData.State == WRITER_FILLING);

+ }

+

+ size_t bytes = MessageSize(*writeMessages.GetVector());

+

+ QuotaReturnSelf(writeMessages.GetVector()->size(), bytes);

+

+ // This is called before `send` syscall inducing latency

+ MessageSent(*writeMessages.GetVector());

+ }

+

+ size_t TRemoteConnection::MessageSize(TArrayRef<TBusMessagePtrAndHeader> messages) {

+ size_t size = 0;

+ for (const auto& message : messages)

+ size += message.MessagePtr->RequestSize;

+

+ return size;

+ }

+

+ size_t TRemoteConnection::GetInFlight() {

+ return AtomicGet(WriterData.InFlight);

+ }

+

+ size_t TRemoteConnection::GetConnectSyscallsNumForTest() {

+ return WriterData.Status.ConnectSyscalls;

+ }

+

+ void TRemoteConnection::WriterBeforeWriteErrorMessage(TBusMessage* message, EMessageStatus status) {

+ if (Session->IsSource_) {

+ CheckedCast<TRemoteClientSession*>(Session.Get())->ReleaseInFlight({message});

+ WriterErrorMessage(message, status);

+ } else {

+ TBusMessagePtrAndHeader h(message);

+ CheckedCast<TRemoteServerSession*>(Session.Get())->ReleaseInWorkResponses(MakeArrayRef(&h, 1));

+ WriterErrorMessage(h.MessagePtr.Release(), status);

+ }

+ }

+

+ void TRemoteConnection::WriterErrorMessage(TNonDestroyingAutoPtr<TBusMessage> m, EMessageStatus status) {

+ TBusMessage* released = m.Release();

+ WriterErrorMessages(MakeArrayRef(&released, 1), status);

+ }

+

+ void TRemoteConnection::WriterErrorMessages(const TArrayRef<TBusMessage*> ms, EMessageStatus status) {

+ ResetOneWayFlag(ms);

+

+ WriterData.Status.Incremental.StatusCounter[status] += ms.size();

+ for (auto m : ms) {

+ Session->InvokeOnError(m, status);

+ }

+ }

+

+ void TRemoteConnection::FireClientConnectionEvent(TClientConnectionEvent::EType type) {

+ Y_VERIFY(Session->IsSource_, "state check");

+ TClientConnectionEvent event(type, ConnectionId, PeerAddr);

+ TRemoteClientSession* session = CheckedCast<TRemoteClientSession*>(Session.Get());

+ session->ClientHandler->OnClientConnectionEvent(event);

+ }

+

+ bool TRemoteConnection::IsAlive() const {

+ return !AtomicGet(WriterData.Down);

+ }

+

+ }

+}