ErrorCode WdtSocket::closeConnectionInternal(bool doTagIOs) {
WVLOG(1) << "Closing socket " << port_ << " " << fd_;
if (fd_ < 0) {
return OK;
}
ErrorCode errorCode = getNonRetryableErrCode();
if (!writesFinalized_) {
errorCode = getMoreInterestingError(errorCode, finalizeWrites(doTagIOs));
}
if (!readsFinalized_) {
errorCode = getMoreInterestingError(errorCode, finalizeReads(doTagIOs));
}
if (::close(fd_) != 0) {
PLOG(ERROR) << "Failed to close socket " << fd_ << " " << port_;
errorCode = getMoreInterestingError(ERROR, errorCode);
}
// This looks like a reset() make it explicit (and check it's complete)
fd_ = -1;
readErrorCode_ = OK;
writeErrorCode_ = OK;
encryptionSettingsRead_ = false;
encryptionSettingsWritten_ = false;
writesFinalized_ = false;
readsFinalized_ = false;
WVLOG(1) << "Error code from close " << errorCodeToStr(errorCode);
return errorCode;
}
void WdtBase::configureThrottler() {
WDT_CHECK(!throttler_);
WVLOG(1) << "Configuring throttler options";
throttler_ = Throttler::makeThrottler(options_);
if (throttler_) {
WLOG(INFO) << "Enabling throttling " << *throttler_;
} else {
WLOG(INFO) << "Throttling not enabled";
}
}
string WdtBase::generateTransferId() {
static std::default_random_engine randomEngine{std::random_device()()};
static std::mutex mutex;
string transferId;
{
std::lock_guard<std::mutex> lock(mutex);
transferId = to_string(randomEngine());
}
WVLOG(1) << "Generated a transfer id " << transferId;
return transferId;
}
void TransferHistoryController::handleGlobalCheckpoint(
const Checkpoint &checkpoint) {
auto errPort = checkpoint.port;
auto it = threadHistoriesMap_.find(errPort);
if (it == threadHistoriesMap_.end()) {
WLOG(ERROR) << "Invalid checkpoint " << checkpoint
<< ". No sender thread running on port " << errPort;
return;
}
WVLOG(1) << "received global checkpoint " << checkpoint;
it->second->setGlobalCheckpoint(checkpoint);
}
int64_t WdtSocket::ioWithAbortCheck(F readOrWrite, T tbuf, int64_t numBytes,
int timeoutMs, bool tryFull) {
WDT_CHECK(threadCtx_.getAbortChecker() != nullptr)
<< "abort checker can not be null";
bool checkAbort = (threadCtx_.getOptions().abort_check_interval_millis > 0);
auto startTime = Clock::now();
int64_t doneBytes = 0;
int retries = 0;
while (doneBytes < numBytes) {
const int64_t ret =
readOrWrite(fd_, tbuf + doneBytes, numBytes - doneBytes);
if (ret < 0) {
// error
if (errno != EINTR && errno != EAGAIN) {
PLOG(ERROR) << "non-retryable error encountered during socket io "
<< fd_ << " " << doneBytes << " " << retries;
return (doneBytes > 0 ? doneBytes : ret);
}
} else if (ret == 0) {
// eof
WVLOG(1) << "EOF received during socket io. fd : " << fd_
<< ", finished bytes : " << doneBytes
<< ", retries : " << retries;
return doneBytes;
} else {
// success
doneBytes += ret;
if (!tryFull) {
// do not have to read/write entire data
return doneBytes;
}
}
if (checkAbort && threadCtx_.getAbortChecker()->shouldAbort()) {
WLOG(ERROR) << "transfer aborted during socket io " << fd_ << " "
<< doneBytes << " " << retries;
return (doneBytes > 0 ? doneBytes : -1);
}
if (timeoutMs > 0) {
int duration = durationMillis(Clock::now() - startTime);
if (duration >= timeoutMs) {
WLOG(INFO) << "socket io timed out after " << duration
<< " ms, retries " << retries << " fd " << fd_
<< " doneBytes " << doneBytes;
return (doneBytes > 0 ? doneBytes : -1);
}
}
retries++;
}
WVLOG_IF(1, retries > 1) << "socket io for " << doneBytes << " bytes took "
<< retries << " retries";
return doneBytes;
}
double Throttler::averageThrottler(const Clock::time_point& now) {
std::chrono::duration<double> elapsedDuration = now - startTime_;
double elapsedSeconds = elapsedDuration.count();
if (avgRateBytesPerSec_ <= 0) {
WVLOG(2) << "There is no avg rate limit";
return -1;
}
const double allowedProgressBytes = avgRateBytesPerSec_ * elapsedSeconds;
if (bytesProgress_ > allowedProgressBytes) {
double idealTime = bytesProgress_ / avgRateBytesPerSec_;
const double sleepTimeSeconds = idealTime - elapsedSeconds;
WVLOG(1) << "Throttler : Elapsed " << elapsedSeconds
<< " seconds. Made progress " << bytesProgress_ / kMbToB
<< " Mbytes in " << elapsedSeconds
<< " seconds, maximum allowed progress for this duration is "
<< allowedProgressBytes / kMbToB
<< " Mbytes. Mean Rate allowed is " << avgRateBytesPerSec_ / kMbToB
<< " Mbytes/sec. Sleeping for " << sleepTimeSeconds << " seconds";
return sleepTimeSeconds;
}
return -1;
}
bool ThreadTransferHistory::addSource(std::unique_ptr<ByteSource> &source) {
std::lock_guard<std::mutex> lock(mutex_);
if (globalCheckpoint_) {
// already received an error for this thread
WVLOG(1) << "adding source after global checkpoint is received. returning "
"the source to the queue";
markSourceAsFailed(source, lastCheckpoint_.get());
lastCheckpoint_.reset();
queue_.returnToQueue(source);
return false;
}
history_.emplace_back(std::move(source));
return true;
}
std::unique_ptr<TransferReport> Receiver::getTransferReport() {
TransferStats globalStats;
for (const auto &receiverThread : receiverThreads_) {
globalStats += receiverThread->getTransferStats();
}
std::unique_ptr<TransferReport> transferReport =
folly::make_unique<TransferReport>(std::move(globalStats));
TransferStatus status = getTransferStatus();
ErrorCode errCode = transferReport->getSummary().getErrorCode();
if (status == NOT_STARTED && errCode == OK) {
WLOG(INFO) << "Transfer not started, setting the error code to ERROR";
transferReport->setErrorCode(ERROR);
}
WVLOG(1) << "Summary code " << errCode;
return transferReport;
}
ErrorCode WdtSocket::finalizeWrites(bool doTagIOs) {
WVLOG(1) << "Finalizing writes/encryption " << port_ << " " << fd_;
ErrorCode code = OK;
std::string tag;
if (!encryptor_.finish(tag)) {
code = ENCRYPTION_ERROR;
}
if (!tag.empty() && doTagIOs) {
const int timeoutMs = threadCtx_.getOptions().write_timeout_millis;
const int expected = tag.size();
if (writeInternal(tag.data(), tag.size(), timeoutMs, false) != expected) {
PLOG(ERROR) << "Encryption Tag write error";
code = ENCRYPTION_ERROR;
}
}
writesFinalized_ = true;
return code;
}
void Receiver::traverseDestinationDir(
std::vector<FileChunksInfo> &fileChunksInfo) {
DirectorySourceQueue dirQueue(options_, destDir_, &abortCheckerCallback_);
dirQueue.buildQueueSynchronously();
auto &discoveredFilesInfo = dirQueue.getDiscoveredFilesMetaData();
for (auto &fileInfo : discoveredFilesInfo) {
if (fileInfo->relPath == kWdtLogName ||
fileInfo->relPath == kWdtBuggyLogName) {
// do not include wdt log files
WVLOG(1) << "Removing " << fileInfo->relPath
<< " from the list of existing files";
continue;
}
FileChunksInfo chunkInfo(fileInfo->seqId, fileInfo->relPath,
fileInfo->size);
chunkInfo.addChunk(Interval(0, fileInfo->size));
fileChunksInfo.emplace_back(std::move(chunkInfo));
}
return;
}
ErrorCode WdtSocket::finalizeReads(bool doTagIOs) {
WVLOG(1) << "Finalizing reads/encryption " << port_ << " " << fd_;
const int toRead = encryptionTypeToTagLen(encryptionParams_.getType());
std::string tag;
ErrorCode code = OK;
if (toRead && doTagIOs) {
tag.resize(toRead);
int read = readInternal(&(tag.front()), tag.size(),
threadCtx_.getOptions().read_timeout_millis, true);
if (read != toRead) {
WLOG(ERROR) << "Unable to read tag at end of stream got " << read
<< " needed " << toRead << " " << folly::humanify(tag);
tag.clear();
code = ENCRYPTION_ERROR;
}
}
if (!decryptor_.finish(tag)) {
code = ENCRYPTION_ERROR;
}
ctxSaveOffset_ = OFFSET_NOT_SET;
readsFinalized_ = true;
return code;
}
double Throttler::calculateSleep(double deltaProgress,
const Clock::time_point& now) {
folly::SpinLockGuard lock(throttlerMutex_);
if (refCount_ <= 0) {
WLOG(ERROR) << "Using the throttler without registering the transfer";
return -1;
}
bytesProgress_ += deltaProgress;
double avgThrottlerSleep = averageThrottler(now);
const bool willSleep = (avgThrottlerSleep > 0);
if (willSleep) {
return avgThrottlerSleep;
}
// we still hold the lock if peak throttler can come into effect
if ((bucketRateBytesPerSec_ > 0) && (bytesTokenBucketLimit_ > 0)) {
std::chrono::duration<double> elapsedDuration = now - lastFillTime_;
lastFillTime_ = now;
double elapsedSeconds = elapsedDuration.count();
bytesTokenBucket_ += elapsedSeconds * bucketRateBytesPerSec_;
if (bytesTokenBucket_ > bytesTokenBucketLimit_) {
bytesTokenBucket_ = bytesTokenBucketLimit_;
}
bytesTokenBucket_ -= deltaProgress;
if (bytesTokenBucket_ < 0) {
/*
* If we have negative number of tokens lets sleep
* This way we will have positive number of tokens next time
*/
double peakThrottlerSleep =
-1.0 * bytesTokenBucket_ / bucketRateBytesPerSec_;
WVLOG(2) << "Peak throttler wants to sleep " << peakThrottlerSleep
<< " seconds";
return peakThrottlerSleep;
}
}
return -1;
}
void WdtBase::setThrottler(std::shared_ptr<Throttler> throttler) {
WVLOG(2) << "Setting an external throttler";
throttler_ = throttler;
}
void TransferHistoryController::addThreadHistory(int32_t port,
TransferStats &threadStats) {
WVLOG(1) << "Adding the history for " << port;
threadHistoriesMap_.emplace(port, std::make_unique<ThreadTransferHistory>(
dirQueue_, threadStats, port));
}
ThreadTransferHistory::ThreadTransferHistory(DirectorySourceQueue &queue,
TransferStats &threadStats,
int32_t port)
: queue_(queue), threadStats_(threadStats), port_(port) {
WVLOG(1) << "Making thread history for port " << port_;
}
WdtSocket::~WdtSocket() {
WVLOG(1) << "~WdtSocket " << port_ << " " << fd_;
closeNoCheck();
}