TySS-Dev/nearby
1
0
Fork 0
mirror of https://github.com/kidfromjupiter/nearby.git synced 2026-09-14 22:56:12 -04:00
Code Issues Packages Projects Releases Wiki Activity

Protect PendingPayload from use-after-free errors.

Browse Source 
Protects against race conditions where a PendingPayload is destroyed when
another thread is still accessing it.

PiperOrigin-RevId: 550990038
This commit is contained in:
Janusz Sobczak
2023-07-25 14:04:11 -07:00
committed by Copybara-Service
parent 5b605b2426
commit ae69b8ff4e
3 changed files with 288 additions and 165 deletions
Download Patch File Download Diff File Expand all files Collapse all files
connections/implementation/BUILD
+2
View File
@@ -137,7 +137,9 @@ cc_library(
"@com_google_absl//absl/container:btree",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/functional:any_invocable",
"@com_google_absl//absl/functional:bind_front",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
connections/implementation/payload_manager.cc
+194 -137
View File
@@ -22,8 +22,11 @@
#include <utility>
#include <vector>
#include "absl/functional/any_invocable.h"
#include "absl/functional/bind_front.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/time/time.h"
#include "connections/implementation/analytics/throughput_recorder.h"
#include "connections/implementation/flags/nearby_connections_feature_flags.h"
@@ -279,9 +282,11 @@ Payload::Id PayloadManager::CreateOutgoingPayload(
NEARBY_LOGS(INFO) << "CreateOutgoingPayload: payload_id=" << payload_id;
MutexLock lock(&mutex_);
pending_payloads_.StartTrackingPayload(
payload_id, absl::make_unique<PendingPayload>(std::move(internal_payload),
endpoint_ids,
/*is_incoming=*/false));
payload_id,
std::make_unique<PendingPayload>(
std::move(internal_payload), endpoint_ids,
/*is_incoming=*/false,
absl::bind_front(&PayloadManager::OnPendingPayloadDestroy, this)));
return payload_id;
}
@@ -297,18 +302,17 @@ void PayloadManager::CancelAllPayloads() {
{
MutexLock lock(&mutex_);
int pending_outgoing_payloads = 0;
for (const auto& pending_id : pending_payloads_.GetAllPayloads()) {
auto* pending = pending_payloads_.GetPayload(pending_id);
pending_payloads_.ForEachPayload([&](PendingPayload* pending) {
if (!pending->IsIncoming()) pending_outgoing_payloads++;
pending->MarkLocallyCanceled();
pending->Close(); // To unblock the sender thread, if there is no data.
}
});
if (pending_outgoing_payloads) {
shutdown_barrier_ =
absl::make_unique<CountDownLatch>(pending_outgoing_payloads);
}
}
if (shutdown_barrier_) {
NEARBY_LOG(INFO,
"PayloadManager: waiting for pending outgoing payloads; self=%p",
@@ -342,9 +346,7 @@ PayloadManager::~PayloadManager() {
NEARBY_LOG(INFO, "PayloadManager: stop tracking payloads; self=%p",
this);
MutexLock lock(&mutex_);
for (const auto& pending_id : pending_payloads_.GetAllPayloads()) {
pending_payloads_.StopTrackingPayload(pending_id);
}
pending_payloads_.StopTrackingAllPayloads();
stop_latch.CountDown();
});
stop_latch.Await();
@@ -419,7 +421,7 @@ void PayloadManager::SendPayload(ClientProxy* client,
"send-payload", [this, client, endpoint_ids, payload_id, payload_type,
resume_offset, payload_total_size]() {
if (shutdown_.Get()) return;
PendingPayload* pending_payload = GetPayload(payload_id);
PendingPayloadHandle pending_payload = GetPayload(payload_id);
if (!pending_payload) {
RecordInvalidPayloadAnalytics(client, endpoint_ids, payload_id,
payload_type, resume_offset,
@@ -468,15 +470,14 @@ void PayloadManager::SendPayload(ClientProxy* client,
<< ", payload_type=" << ToString(payload_type);
}
PayloadManager::PendingPayload* PayloadManager::GetPayload(
PayloadManager::PendingPayloadHandle PayloadManager::GetPayload(
Payload::Id payload_id) const {
MutexLock lock(&mutex_);
return pending_payloads_.GetPayload(payload_id);
}
Status PayloadManager::CancelPayload(ClientProxy* client,
Payload::Id payload_id) {
PendingPayload* canceled_payload = GetPayload(payload_id);
PendingPayloadHandle canceled_payload = GetPayload(payload_id);
if (!canceled_payload) {
NEARBY_LOGS(INFO) << "Client requested cancel for unknown payload_id="
<< payload_id << ", ignoring.";
@@ -532,51 +533,49 @@ void PayloadManager::OnEndpointDisconnect(ClientProxy* client,
}
RunOnStatusUpdateThread(
"payload-manager-on-disconnect",
[this, client, endpoint_id, barrier]()
RUN_ON_PAYLOAD_STATUS_UPDATE_THREAD() mutable {
// Iterate through all our payloads and look for payloads associated
// with this endpoint.
MutexLock lock(&mutex_);
for (const auto& payload_id : pending_payloads_.GetAllPayloads()) {
auto* pending_payload = pending_payloads_.GetPayload(payload_id);
if (!pending_payload) continue;
auto endpoint_info = pending_payload->GetEndpoint(endpoint_id);
if (!endpoint_info) continue;
std::int64_t endpoint_offset = endpoint_info->offset;
// Stop tracking the endpoint for this payload.
pending_payload->RemoveEndpoints({endpoint_id});
// |endpoint_info| is longer valid after calling RemoveEndpoints.
endpoint_info = nullptr;
[this, client, endpoint_id,
barrier]() RUN_ON_PAYLOAD_STATUS_UPDATE_THREAD() mutable {
// Iterate through all our payloads and look for payloads associated
// with this endpoint.
MutexLock lock(&mutex_);
pending_payloads_.ForEachPayload([&](PendingPayload* pending_payload) {
auto endpoint_info = pending_payload->GetEndpoint(endpoint_id);
if (!endpoint_info) return;
std::int64_t endpoint_offset = endpoint_info->offset;
// Stop tracking the endpoint for this payload.
pending_payload->RemoveEndpoints({endpoint_id});
// |endpoint_info| is longer valid after calling
// RemoveEndpoints.
endpoint_info = nullptr;
std::int64_t payload_total_size =
pending_payload->GetInternalPayload()->GetTotalSize();
std::int64_t payload_total_size =
pending_payload->GetInternalPayload()->GetTotalSize();
// If no endpoints are left for this payload, close it.
if (pending_payload->GetEndpoints().empty()) {
pending_payload->Close();
}
// If no endpoints are left for this payload, close it.
if (pending_payload->GetEndpoints().empty()) {
pending_payload->Close();
}
// Create the payload transfer update.
PayloadProgressInfo update{pending_payload->GetId(),
PayloadProgressInfo::Status::kFailure,
payload_total_size, endpoint_offset};
// Create the payload transfer update.
PayloadProgressInfo update{payload_id,
PayloadProgressInfo::Status::kFailure,
payload_total_size, endpoint_offset};
// Send a client notification of a payload transfer failure.
client->OnPayloadProgress(endpoint_id, update);
// Send a client notification of a payload transfer failure.
client->OnPayloadProgress(endpoint_id, update);
if (pending_payload->IsIncoming()) {
client->GetAnalyticsRecorder().OnIncomingPayloadDone(
endpoint_id, pending_payload->GetId(),
location::nearby::proto::connections::ENDPOINT_IO_ERROR);
} else {
client->GetAnalyticsRecorder().OnOutgoingPayloadDone(
endpoint_id, pending_payload->GetId(),
location::nearby::proto::connections::ENDPOINT_IO_ERROR);
}
});
if (pending_payload->IsIncoming()) {
client->GetAnalyticsRecorder().OnIncomingPayloadDone(
endpoint_id, pending_payload->GetId(),
location::nearby::proto::connections::ENDPOINT_IO_ERROR);
} else {
client->GetAnalyticsRecorder().OnOutgoingPayloadDone(
endpoint_id, pending_payload->GetId(),
location::nearby::proto::connections::ENDPOINT_IO_ERROR);
}
}
barrier.CountDown();
});
barrier.CountDown();
});
}
location::nearby::proto::connections::PayloadStatus
@@ -686,25 +685,33 @@ PayloadTransferFrame::PayloadChunk PayloadManager::CreatePayloadChunk(
return payload_chunk;
}
PayloadManager::PendingPayload* PayloadManager::CreateIncomingPayload(
PayloadManager::PendingPayloadHandle PayloadManager::CreateIncomingPayload(
const PayloadTransferFrame& frame, const std::string& endpoint_id) {
auto internal_payload =
CreateIncomingInternalPayload(frame, custom_save_path_);
if (!internal_payload) {
return nullptr;
return PendingPayloadHandle();
}
Payload::Id payload_id = internal_payload->GetId();
NEARBY_LOGS(INFO) << "CreateIncomingPayload: payload_id=" << payload_id;
MutexLock lock(&mutex_);
pending_payloads_.StartTrackingPayload(
payload_id,
absl::make_unique<PendingPayload>(std::move(internal_payload),
EndpointIds{endpoint_id}, true));
std::make_unique<PendingPayload>(
std::move(internal_payload), EndpointIds{endpoint_id}, true,
absl::bind_front(&PayloadManager::OnPendingPayloadDestroy, this)));
return pending_payloads_.GetPayload(payload_id);
}
void PayloadManager::OnPendingPayloadDestroy(const PendingPayload* payload) {
NEARBY_LOGS(INFO) << "PayloadManager: destroying " << payload->ToString()
<< " self=" << this;
if (payload->IsIncoming()) return;
RunOnStatusUpdateThread(
"~PendingPayload",
[this]() RUN_ON_PAYLOAD_STATUS_UPDATE_THREAD() { NotifyShutdown(); });
}
void PayloadManager::SendClientCallbacksForFinishedOutgoingPayload(
ClientProxy* client, const EndpointIds& finished_endpoint_ids,
const PayloadTransferFrame::PayloadHeader& payload_header,
@@ -716,7 +723,7 @@ void PayloadManager::SendClientCallbacksForFinishedOutgoingPayload(
num_bytes_successfully_transferred,
status]() RUN_ON_PAYLOAD_STATUS_UPDATE_THREAD() {
// Make sure we're still tracking this payload.
PendingPayload* pending_payload = GetPayload(payload_header.id());
PendingPayloadHandle pending_payload = GetPayload(payload_header.id());
if (!pending_payload) {
return;
}
@@ -760,13 +767,14 @@ void PayloadManager::SendClientCallbacksForFinishedIncomingPayload(
[this, client, endpoint_id, payload_header, offset_bytes,
status]() RUN_ON_PAYLOAD_STATUS_UPDATE_THREAD() {
// Make sure we're still tracking this payload.
PendingPayload* pending_payload = GetPayload(payload_header.id());
PendingPayloadHandle pending_payload = GetPayload(payload_header.id());
if (!pending_payload) {
return;
}
// Unless we never started tracking this payload (meaning we failed to
// even create the InternalPayload), notify the client (and close it).
// Unless we never started tracking this payload (meaning we
// failed to even create the InternalPayload), notify the client
// (and close it).
PayloadProgressInfo update{
payload_header.id(),
PayloadManager::PayloadStatusToTransferUpdateStatus(status),
@@ -908,7 +916,7 @@ void PayloadManager::HandleSuccessfulOutgoingChunk(
}
}
PendingPayload* pending_payload = GetPayload(payload_header.id());
PendingPayloadHandle pending_payload = GetPayload(payload_header.id());
if (!pending_payload || !pending_payload->GetEndpoint(endpoint_id)) {
NEARBY_LOGS(INFO)
<< "HandleSuccessfulOutgoingChunk: endpoint not found: "
@@ -949,20 +957,7 @@ void PayloadManager::HandleSuccessfulOutgoingChunk(
// @PayloadManagerStatusUpdateThread
void PayloadManager::DestroyPendingPayload(Payload::Id payload_id) {
bool is_incoming = false;
{
MutexLock lock(&mutex_);
auto pending = pending_payloads_.StopTrackingPayload(payload_id);
if (!pending) return;
is_incoming = pending->IsIncoming();
const char* direction = is_incoming ? "incoming" : "outgoing";
NEARBY_LOGS(INFO) << "PayloadManager: destroying " << direction
<< " pending payload: self=" << this
<< "; payload_id=" << payload_id;
pending->Close();
pending.reset();
}
if (!is_incoming) NotifyShutdown();
pending_payloads_.StopTrackingPayload(payload_id);
}
void PayloadManager::HandleSuccessfulIncomingChunk(
@@ -1003,7 +998,7 @@ void PayloadManager::HandleSuccessfulIncomingChunk(
}
}
PendingPayload* pending_payload = GetPayload(payload_header.id());
PendingPayloadHandle pending_payload = GetPayload(payload_header.id());
if (!pending_payload) {
return;
}
@@ -1044,11 +1039,11 @@ void PayloadManager::ProcessDataPacket(
<< payload_header.id()
<< " from endpoint_id=" << from_endpoint_id
<< " at offset " << payload_chunk.offset();
PendingPayload* pending_payload;
Payload::Id payload_id = payload_header.id();
PendingPayloadHandle pending_payload;
if (payload_chunk.offset() == 0) {
ThroughputRecorderContainer::GetInstance()
.GetTPRecorder(payload_header.id(), PayloadDirection::INCOMING_PAYLOAD)
.GetTPRecorder(payload_id, PayloadDirection::INCOMING_PAYLOAD)
->Start((PayloadType)payload_header.type(),
PayloadDirection::INCOMING_PAYLOAD);
packet_meta_data.Reset();
@@ -1077,12 +1072,13 @@ void PayloadManager::ProcessDataPacket(
PayloadTransferFrame::ControlMessage::PAYLOAD_ERROR);
return;
}
// Also, let the client know of this new incoming payload.
RunOnStatusUpdateThread(
"process-data-packet",
[to_client, from_endpoint_id, pending_payload]()
[to_client, from_endpoint_id,
pending_payload = GetPayload(payload_id)]()
RUN_ON_PAYLOAD_STATUS_UPDATE_THREAD() {
if (!pending_payload) return;
NEARBY_LOGS(INFO)
<< "PayloadManager received new payload_id="
<< pending_payload->GetInternalPayload()->GetId()
@@ -1093,17 +1089,17 @@ void PayloadManager::ProcessDataPacket(
});
} else {
pending_payload = GetPayload(payload_header.id());
if (!pending_payload) {
NEARBY_LOGS(WARNING) << "ProcessDataPacket: [missing] endpoint_id="
<< from_endpoint_id
<< "; payload_id=" << payload_header.id();
return;
}
}
if (!pending_payload) {
NEARBY_LOGS(WARNING) << "ProcessDataPacket: [missing] endpoint_id="
<< from_endpoint_id
<< "; payload_id=" << payload_header.id();
return;
}
if (pending_payload->IsLocallyCanceled()) {
// This incoming payload was canceled by the client. Drop this frame and do
// all the cleanup. See go/nc-cancel-payload
// This incoming payload was canceled by the client. Drop this frame and
// do all the cleanup. See go/nc-cancel-payload
NEARBY_LOGS(INFO) << "ProcessDataPacket: [cancel] endpoint_id="
<< from_endpoint_id
<< "; payload_id=" << pending_payload->GetId();
@@ -1115,10 +1111,10 @@ void PayloadManager::ProcessDataPacket(
}
// Update the offset for this payload. An endpoint disconnection might occur
// from another thread and we would need to know the current offset to report
// back to the client. For the sake of accuracy, we update the pending payload
// here because it's after all payload terminating events are handled, but
// right before we actually start attaching the next chunk.
// from another thread and we would need to know the current offset to
// report back to the client. For the sake of accuracy, we update the
// pending payload here because it's after all payload terminating events
// are handled, but right before we actually start attaching the next chunk.
pending_payload->SetOffsetForEndpoint(from_endpoint_id,
payload_chunk.offset());
@@ -1165,7 +1161,7 @@ void PayloadManager::ProcessControlPacket(
payload_transfer_frame.payload_header();
const PayloadTransferFrame::ControlMessage& control_message =
payload_transfer_frame.control_message();
PendingPayload* pending_payload = GetPayload(payload_header.id());
PendingPayloadHandle pending_payload = GetPayload(payload_header.id());
if (!pending_payload) {
NEARBY_LOGS(INFO) << "Got ControlMessage for unknown payload_id="
<< payload_header.id()
@@ -1266,7 +1262,8 @@ void PayloadManager::SetCustomSavePath(ClientProxy* client,
custom_save_path_ = path;
}
///////////////////////////////// EndpointInfo /////////////////////////////////
///////////////////////////////// EndpointInfo
////////////////////////////////////
PayloadManager::EndpointInfo::Status
PayloadManager::EndpointInfo::ControlMessageEventToEndpointInfoStatus(
@@ -1292,13 +1289,16 @@ void PayloadManager::EndpointInfo::SetStatusFromControlMessage(
<< " based on OOB ControlMessage";
}
//////////////////////////////// PendingPayload ////////////////////////////////
//////////////////////////////// PendingPayload
///////////////////////////////////
PayloadManager::PendingPayload::PendingPayload(
std::unique_ptr<InternalPayload> internal_payload,
const EndpointIds& endpoint_ids, bool is_incoming)
const EndpointIds& endpoint_ids, bool is_incoming,
DestroyCallback destroy_callback)
: is_incoming_(is_incoming),
internal_payload_(std::move(internal_payload)) {
internal_payload_(std::move(internal_payload)),
destroy_callback_(std::move(destroy_callback)) {
// Initially we mark all endpoints as available.
// Later on some may become canceled, some may experience data transfer
// failures. Any of these situations will cause endpoint to be marked as
@@ -1384,66 +1384,123 @@ void PayloadManager::PendingPayload::SetOffsetForEndpoint(
}
void PayloadManager::PendingPayload::Close() {
bool was_closed = is_closed_.Set(true);
if (was_closed) return;
if (internal_payload_) internal_payload_->Close();
close_event_.CountDown();
}
bool PayloadManager::PendingPayload::WaitForClose() {
return close_event_.Await(kWaitCloseTimeout).result();
}
bool PayloadManager::PendingPayload::IsClosed() {
return close_event_.Await(absl::ZeroDuration()).result();
}
void PayloadManager::RunOnStatusUpdateThread(const std::string& name,
std::function<void()> runnable) {
void PayloadManager::RunOnStatusUpdateThread(
const std::string& name, absl::AnyInvocable<void()> runnable) {
payload_status_update_executor_.Execute(name, std::move(runnable));
}
/////////////////////////////// PendingPayloads ///////////////////////////////
/////////////////////////////// PendingPayloads
//////////////////////////////////
void PayloadManager::PendingPayloads::StartTrackingPayload(
Payload::Id payload_id, std::unique_ptr<PendingPayload> pending_payload) {
MutexLock lock(&mutex_);
// If the |payload_id| is being re-used, always prefer the newer payload.
auto it = pending_payloads_.find(payload_id);
if (it != pending_payloads_.end()) {
pending_payloads_.erase(payload_id);
}
auto pair = pending_payloads_.emplace(payload_id, std::move(pending_payload));
NEARBY_LOGS(INFO) << "StartTrackingPayload: payload_id=" << payload_id
<< "; inserted=" << pair.second;
Remove(pending_payloads_.find(payload_id));
NEARBY_LOGS(INFO) << "StartTrackingPayload: " << pending_payload->ToString();
pending_payload->IncRefCount();
pending_payloads_[payload_id] = std::move(pending_payload);
}
std::unique_ptr<PayloadManager::PendingPayload>
PayloadManager::PendingPayloads::StopTrackingPayload(Payload::Id payload_id) {
void PayloadManager::PendingPayloads::StopTrackingPayload(
Payload::Id payload_id) {
MutexLock lock(&mutex_);
auto it = pending_payloads_.find(payload_id);
if (it == pending_payloads_.end()) return {};
auto item = pending_payloads_.extract(it);
return std::move(item.mapped());
NEARBY_LOGS(INFO) << "StopTrackingPayload " << payload_id;
Remove(pending_payloads_.find(payload_id));
}
PayloadManager::PendingPayload* PayloadManager::PendingPayloads::GetPayload(
Payload::Id payload_id) const {
void PayloadManager::PendingPayloads::Remove(
absl::flat_hash_map<Payload::Id, std::unique_ptr<PendingPayload>>::iterator
it) {
if (it != pending_payloads_.end()) {
int refcount = it->second->DecRefCount();
if (refcount == 0) {
// Nobody is using the payload, we can remove it.
NEARBY_LOGS(VERBOSE) << "Erase payload " << it->second->ToString();
pending_payloads_.erase(it);
} else {
// Someone is still using the payload. Move it to the garbage bin. The
// payload will be removed when they release it.
NEARBY_LOGS(VERBOSE) << "Bin payload " << it->second->ToString();
payload_garbage_bin_.push_back(
std::move(pending_payloads_.extract(it).mapped()));
}
}
}
PayloadManager::PendingPayloadHandle
PayloadManager::PendingPayloads::GetPayload(Payload::Id payload_id) const {
MutexLock lock(&mutex_);
auto item = pending_payloads_.find(payload_id);
return item != pending_payloads_.end() ? item->second.get() : nullptr;
if (item == pending_payloads_.end()) {
return PendingPayloadHandle();
}
PendingPayload* payload = item->second.get();
payload->IncRefCount();
return PendingPayloadHandle(
payload, absl::bind_front(&PendingPayloads::Release,
const_cast<PendingPayloads*>(this)));
}
std::vector<Payload::Id> PayloadManager::PendingPayloads::GetAllPayloads() {
void PayloadManager::PendingPayloads::StopTrackingAllPayloads() {
MutexLock lock(&mutex_);
std::vector<Payload::Id> result;
for (const auto& item : pending_payloads_) {
result.push_back(item.first);
for (auto it = pending_payloads_.begin(); it != pending_payloads_.end();) {
Remove(it++);
}
return result;
}
void PayloadManager::PendingPayloads::ForEachPayload(
absl::AnyInvocable<void(PendingPayload*)> callback) {
MutexLock lock(&mutex_);
for (const auto& item : pending_payloads_) {
callback(item.second.get());
}
}
void PayloadManager::PendingPayloads::Release(PendingPayload* payload) {
// Called when `PendingPayloadHandle` is destroyed.
MutexLock lock(&mutex_);
NEARBY_LOGS(VERBOSE) << __func__ << " " << payload->ToString();
auto it = pending_payloads_.find(payload->GetId());
if (it != pending_payloads_.end() && it->second.get() == payload) {
// The payload is still tracked.
payload->DecRefCount();
return;
}
auto bin_it =
std::find_if(payload_garbage_bin_.begin(), payload_garbage_bin_.end(),
[payload](auto& item) { return item.get() == payload; });
if (bin_it != payload_garbage_bin_.end()) {
int refcount = payload->DecRefCount();
if (refcount == 0) {
// The payload is not tracked and it was the last reference.
payload_garbage_bin_.erase(bin_it);
}
}
}
PayloadManager::PendingPayloadHandle::PendingPayloadHandle(
PendingPayload* payload, DestroyCallback destroy_callback)
: payload_(payload), destroy_callback_(std::move(destroy_callback)) {}
PayloadManager::PendingPayloadHandle::~PendingPayloadHandle() {
if (destroy_callback_) {
std::move(destroy_callback_)(payload_);
}
}
std::string PayloadManager::PendingPayload::ToString() const {
return absl::StrFormat("Payload(%s, %d)",
IsIncoming() ? "incoming" : "outgoing", GetId());
}
} // namespace connections
connections/implementation/payload_manager.h
+92 -28
View File
@@ -15,6 +15,7 @@
#ifndef CORE_INTERNAL_PAYLOAD_MANAGER_H_
#define CORE_INTERNAL_PAYLOAD_MANAGER_H_
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
@@ -23,6 +24,7 @@
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/functional/any_invocable.h"
#include "connections/implementation/analytics/packet_meta_data.h"
#include "connections/implementation/client_proxy.h"
#include "connections/implementation/endpoint_manager.h"
@@ -99,12 +101,19 @@ class PayloadManager : public EndpointManager::FrameProcessor {
// Tracks state for an InternalPayload and the endpoints associated with it.
class PendingPayload {
public:
using DestroyCallback = absl::AnyInvocable<void(PendingPayload*) &&>;
PendingPayload(std::unique_ptr<InternalPayload> internal_payload,
const EndpointIds& endpoint_ids, bool is_incoming);
const EndpointIds& endpoint_ids, bool is_incoming,
DestroyCallback destroy_callback);
PendingPayload(PendingPayload&&) = default;
PendingPayload& operator=(PendingPayload&&) = default;
~PendingPayload() { Close(); }
~PendingPayload() {
Close();
if (destroy_callback_) {
std::move(destroy_callback_)(this);
}
}
Payload::Id GetId() const;
@@ -137,24 +146,65 @@ class PayloadManager : public EndpointManager::FrameProcessor {
void SetOffsetForEndpoint(const std::string& endpoint_id,
std::int64_t offset) ABSL_LOCKS_EXCLUDED(mutex_);
// Closes internal_payload_ and triggers close_event_.
// Closes internal_payload_.
// Close is called when a pending peyload does not have associated
// endpoints.
void Close();
// Waits for close_event_ or for timeout to happen.
// Returns true, if event happened, false otherwise.
bool WaitForClose();
bool IsClosed();
std::string ToString() const;
// Ref counting for `PendingPayloads` use only. `PendingPayloads` class owns
// all instances of `PendingPayload`.
int IncRefCount() { return ++refcount_; }
int DecRefCount() { return --refcount_; }
private:
mutable Mutex mutex_;
bool is_incoming_;
AtomicBoolean is_locally_canceled_{false};
CountDownLatch close_event_{1};
AtomicBoolean is_closed_;
std::unique_ptr<InternalPayload> internal_payload_;
DestroyCallback destroy_callback_;
absl::flat_hash_map<std::string, EndpointInfo> endpoints_
ABSL_GUARDED_BY(mutex_);
int refcount_ = 0;
};
// A RAII handle to `PendingPayload`. Holding a `PendingPayloadHandle`
// guarantees that `PendingPaylaod` won't be destroyed while in use.
// Create instances with `GetPayload(Payload::Id)`.
class PendingPayloadHandle {
public:
using DestroyCallback = absl::AnyInvocable<void(PendingPayload*) &&>;
PendingPayloadHandle() = default;
PendingPayloadHandle(PendingPayload* payload,
DestroyCallback destroy_callback);
PendingPayloadHandle(const PendingPayloadHandle&) = delete;
PendingPayloadHandle(PendingPayloadHandle&& other) {
payload_ = other.payload_;
other.payload_ = nullptr;
destroy_callback_ = std::move(other.destroy_callback_);
}
~PendingPayloadHandle();
PendingPayloadHandle& operator=(const PendingPayloadHandle&) = delete;
PendingPayloadHandle& operator=(PendingPayloadHandle&& other) {
if (payload_ != nullptr && destroy_callback_) {
std::move(destroy_callback_)(payload_);
}
payload_ = other.payload_;
other.payload_ = nullptr;
destroy_callback_ = std::move(other.destroy_callback_);
return *this;
}
explicit operator bool() const { return payload_ != nullptr; }
PendingPayload* operator->() const { return payload_; }
PendingPayload& operator*() const { return *payload_; }
private:
PendingPayload* payload_ = nullptr;
DestroyCallback destroy_callback_;
};
// Tracks and manages PendingPayload objects in a synchronized manner.
@@ -166,16 +216,30 @@ class PayloadManager : public EndpointManager::FrameProcessor {
void StartTrackingPayload(Payload::Id payload_id,
std::unique_ptr<PendingPayload> pending_payload)
ABSL_LOCKS_EXCLUDED(mutex_);
std::unique_ptr<PendingPayload> StopTrackingPayload(Payload::Id payload_id)
void StopTrackingPayload(Payload::Id payload_id)
ABSL_LOCKS_EXCLUDED(mutex_);
PendingPayload* GetPayload(Payload::Id payload_id) const
void StopTrackingAllPayloads() ABSL_LOCKS_EXCLUDED(mutex_);
PendingPayloadHandle GetPayload(Payload::Id payload_id) const
ABSL_LOCKS_EXCLUDED(mutex_);
// Calls `callback` for each tracked payload. The callback must not call
// other `PendingPayloads` methods.
void ForEachPayload(absl::AnyInvocable<void(PendingPayload*)> callback)
ABSL_LOCKS_EXCLUDED(mutex_);
std::vector<Payload::Id> GetAllPayloads() ABSL_LOCKS_EXCLUDED(mutex_);
private:
void Release(PendingPayload* payload) ABSL_LOCKS_EXCLUDED(mutex_);
void Remove(absl::flat_hash_map<
Payload::Id, std::unique_ptr<PendingPayload>>::iterator it)
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
mutable Mutex mutex_;
absl::flat_hash_map<Payload::Id, std::unique_ptr<PendingPayload>>
pending_payloads_ ABSL_GUARDED_BY(mutex_);
// When we stop tracking a payload but someone is still holding a handle to
// the payload, we can't delete it just yet. Instead, we move it to the
// garbage bin. When the `PendingPayloadHandle` is released, the payload
// will be removed from the bin.
std::vector<std::unique_ptr<PendingPayload>> payload_garbage_bin_
ABSL_GUARDED_BY(mutex_);
};
using Endpoints = std::vector<const EndpointInfo*>;
@@ -185,8 +249,8 @@ class PayloadManager : public EndpointManager::FrameProcessor {
static std::string ToString(EndpointInfo::Status status);
// Splits the endpoints for this payload by availability.
// Returns a pair of lists of EndpointInfo*, with the first being the list of
// still-available endpoints, and the second for unavailable endpoints.
// Returns a pair of lists of EndpointInfo*, with the first being the list
// of still-available endpoints, and the second for unavailable endpoints.
static std::pair<Endpoints, Endpoints> GetAvailableAndUnavailableEndpoints(
const PendingPayload& pending_payload);
@@ -203,14 +267,14 @@ class PayloadManager : public EndpointManager::FrameProcessor {
std::int64_t offset_bytes,
location::nearby::proto::connections::PayloadStatus status);
// Converts the status of an endpoint that's been set out-of-band via a remote
// ControlMessage to the PayloadStatus for handling of that endpoint-payload
// pair.
// Converts the status of an endpoint that's been set out-of-band via a
// remote ControlMessage to the PayloadStatus for handling of that
// endpoint-payload pair.
static location::nearby::proto::connections::PayloadStatus
EndpointInfoStatusToPayloadStatus(EndpointInfo::Status status);
// Converts a ControlMessage::EventType for a particular payload to a
// PayloadStatus. Called when we've received a ControlMessage with this event
// from a remote endpoint; thus the PayloadStatuses are REMOTE_*.
// PayloadStatus. Called when we've received a ControlMessage with this
// event from a remote endpoint; thus the PayloadStatuses are REMOTE_*.
static location::nearby::proto::connections::PayloadStatus
ControlMessageEventToPayloadStatus(
PayloadTransferFrame::ControlMessage::EventType event);
@@ -226,8 +290,8 @@ class PayloadManager : public EndpointManager::FrameProcessor {
PayloadTransferFrame::PayloadChunk CreatePayloadChunk(std::int64_t offset,
ByteArray body);
PendingPayload* CreateIncomingPayload(const PayloadTransferFrame& frame,
const std::string& endpoint_id)
PendingPayloadHandle CreateIncomingPayload(const PayloadTransferFrame& frame,
const std::string& endpoint_id)
ABSL_LOCKS_EXCLUDED(mutex_);
Payload::Id CreateOutgoingPayload(Payload payload,
@@ -251,8 +315,8 @@ class PayloadManager : public EndpointManager::FrameProcessor {
std::int64_t num_bytes_successfully_transferred,
PayloadTransferFrame::ControlMessage::EventType event_type);
// Handles a finished outgoing payload for the given endpointIds. All statuses
// except for SUCCESS are handled here.
// Handles a finished outgoing payload for the given endpointIds. All
// statuses except for SUCCESS are handled here.
void HandleFinishedOutgoingPayload(
ClientProxy* client, const EndpointIds& finished_endpoint_ids,
const PayloadTransferFrame::PayloadHeader& payload_header,
@@ -293,11 +357,11 @@ class PayloadManager : public EndpointManager::FrameProcessor {
SingleThreadExecutor* GetOutgoingPayloadExecutor(PayloadType payload_type);
void RunOnStatusUpdateThread(const std::string& name,
std::function<void()> runnable);
absl::AnyInvocable<void()> runnable);
bool NotifyShutdown() ABSL_LOCKS_EXCLUDED(mutex_);
void DestroyPendingPayload(Payload::Id payload_id)
ABSL_LOCKS_EXCLUDED(mutex_);
PendingPayload* GetPayload(Payload::Id payload_id) const
PendingPayloadHandle GetPayload(Payload::Id payload_id) const
ABSL_LOCKS_EXCLUDED(mutex_);
void CancelAllPayloads() ABSL_LOCKS_EXCLUDED(mutex_);
@@ -317,23 +381,23 @@ class PayloadManager : public EndpointManager::FrameProcessor {
PayloadType FramePayloadTypeToPayloadType(
PayloadTransferFrame::PayloadHeader::PayloadType type);
void OnPendingPayloadDestroy(const PendingPayload* payload);
mutable Mutex mutex_;
std::string custom_save_path_;
AtomicBoolean shutdown_{false};
std::unique_ptr<CountDownLatch> shutdown_barrier_;
int send_payload_count_ = 0;
PendingPayloads pending_payloads_ ABSL_GUARDED_BY(mutex_);
SingleThreadExecutor bytes_payload_executor_;
SingleThreadExecutor file_payload_executor_;
SingleThreadExecutor stream_payload_executor_;
SingleThreadExecutor payload_status_update_executor_;
PendingPayloads pending_payloads_;
EndpointManager* endpoint_manager_;
// When callback processing cannot keep the speed of callback update, the
// callback thread will be lag to the real transfer. In order to keep sync
// between callback and sending/receiving threads, we will skip non-important
// callbacks during file transfer.
// between callback and sending/receiving threads, we will skip
// non-important callbacks during file transfer.
mutable Mutex chunk_update_mutex_;
int outgoing_chunk_update_count_ ABSL_GUARDED_BY(chunk_update_mutex_) = 0;
int incoming_chunk_update_count_ ABSL_GUARDED_BY(chunk_update_mutex_) = 0;