This commit is contained in:
kssenii 2022-04-28 12:57:22 +02:00
parent 37d11cfb1c
commit 9b38c9f7eb
3 changed files with 39 additions and 37 deletions

View File

@ -653,9 +653,7 @@ void FileSegment::assertNotDetached() const
void FileSegment::assertDetachedStatus() const void FileSegment::assertDetachedStatus() const
{ {
assert( assert(download_state == State::EMPTY || hasFinalizedState());
(download_state == State::EMPTY) || (download_state == State::PARTIALLY_DOWNLOADED_NO_CONTINUATION)
|| (download_state == State::SKIP_CACHE));
} }
FileSegmentPtr FileSegment::getSnapshot(const FileSegmentPtr & file_segment, std::lock_guard<std::mutex> & /* cache_lock */) FileSegmentPtr FileSegment::getSnapshot(const FileSegmentPtr & file_segment, std::lock_guard<std::mutex> & /* cache_lock */)
@ -675,18 +673,21 @@ FileSegmentPtr FileSegment::getSnapshot(const FileSegmentPtr & file_segment, std
return snapshot; return snapshot;
} }
bool FileSegment::hasFinalizedState() const
{
return download_state == State::DOWNLOADED
|| download_state == State::PARTIALLY_DOWNLOADED_NO_CONTINUATION
|| download_state == State::SKIP_CACHE;
}
void FileSegment::detach(std::lock_guard<std::mutex> & cache_lock, std::lock_guard<std::mutex> & segment_lock) void FileSegment::detach(std::lock_guard<std::mutex> & cache_lock, std::lock_guard<std::mutex> & segment_lock)
{ {
if (detached) if (detached)
return; return;
bool has_finalized_state = download_state == State::DOWNLOADED
|| download_state == State::PARTIALLY_DOWNLOADED_NO_CONTINUATION
|| download_state == State::SKIP_CACHE;
detached = true; detached = true;
if (!has_finalized_state) if (!hasFinalizedState())
{ {
completeUnlocked(cache_lock, segment_lock); completeUnlocked(cache_lock, segment_lock);
} }

View File

@ -154,7 +154,7 @@ private:
void assertCorrectnessImpl(std::lock_guard<std::mutex> & segment_lock) const; void assertCorrectnessImpl(std::lock_guard<std::mutex> & segment_lock) const;
void assertNotDetached() const; void assertNotDetached() const;
void assertDetachedStatus() const; void assertDetachedStatus() const;
bool hasFinalizedState() const;
void setDownloaded(std::lock_guard<std::mutex> & segment_lock); void setDownloaded(std::lock_guard<std::mutex> & segment_lock);
void setDownloadFailed(std::lock_guard<std::mutex> & segment_lock); void setDownloadFailed(std::lock_guard<std::mutex> & segment_lock);

View File

@ -92,6 +92,35 @@ void WriteBufferFromS3::nextImpl()
? CurrentThread::get().getThreadGroup() ? CurrentThread::get().getThreadGroup()
: MainThreadStatus::getInstance().getThreadGroup(); : MainThreadStatus::getInstance().getThreadGroup();
if (cacheEnabled())
{
auto cache_key = cache->hash(key);
file_segments_holder.emplace(cache->setDownloading(cache_key, current_download_offset, size));
current_download_offset += size;
size_t remaining_size = size;
auto & file_segments = file_segments_holder->file_segments;
for (auto file_segment_it = file_segments.begin(); file_segment_it != file_segments.end(); ++file_segment_it)
{
auto & file_segment = *file_segment_it;
size_t current_size = std::min(file_segment->range().size(), remaining_size);
remaining_size -= current_size;
if (file_segment->reserve(current_size))
{
file_segment->writeInMemory(working_buffer.begin(), current_size);
}
else
{
for (auto reset_segment_it = file_segment_it; reset_segment_it != file_segments.end(); ++reset_segment_it)
(*reset_segment_it)->complete(FileSegment::State::PARTIALLY_DOWNLOADED_NO_CONTINUATION);
file_segments.erase(file_segment_it, file_segments.end());
break;
}
}
}
ProfileEvents::increment(ProfileEvents::S3WriteBytes, offset()); ProfileEvents::increment(ProfileEvents::S3WriteBytes, offset());
last_part_size += offset(); last_part_size += offset();
@ -102,34 +131,6 @@ void WriteBufferFromS3::nextImpl()
if (!multipart_upload_id.empty() && last_part_size > upload_part_size) if (!multipart_upload_id.empty() && last_part_size > upload_part_size)
{ {
if (cacheEnabled())
{
auto cache_key = cache->hash(key);
file_segments_holder.emplace(cache->setDownloading(cache_key, current_download_offset, size));
current_download_offset += size;
size_t remaining_size = size;
auto & file_segments = file_segments_holder->file_segments;
for (auto file_segment_it = file_segments.begin(); file_segment_it != file_segments.end(); ++file_segment_it)
{
auto & file_segment = *file_segment_it;
size_t current_size = std::min(file_segment->range().size(), remaining_size);
remaining_size -= current_size;
if (file_segment->reserve(current_size))
{
file_segment->writeInMemory(working_buffer.begin(), current_size);
}
else
{
for (auto reset_segment_it = file_segment_it; reset_segment_it != file_segments.end(); ++reset_segment_it)
(*reset_segment_it)->complete(FileSegment::State::PARTIALLY_DOWNLOADED_NO_CONTINUATION);
file_segments.erase(file_segment_it, file_segments.end());
break;
}
}
}
writePart(); writePart();
allocateBuffer(); allocateBuffer();