diff options
Diffstat (limited to 'src/zenutil/filebuildstorage.cpp')
| -rw-r--r-- | src/zenutil/filebuildstorage.cpp | 26 |
1 files changed, 16 insertions, 10 deletions
diff --git a/src/zenutil/filebuildstorage.cpp b/src/zenutil/filebuildstorage.cpp index 7aa252e44..f335a03a3 100644 --- a/src/zenutil/filebuildstorage.cpp +++ b/src/zenutil/filebuildstorage.cpp @@ -369,11 +369,11 @@ public: return IoBuffer{}; } - virtual std::vector<std::function<void()>> GetLargeBuildBlob( - const Oid& BuildId, - const IoHash& RawHash, - uint64_t ChunkSize, - std::function<void(uint64_t Offset, const IoBuffer& Chunk, uint64_t BytesRemaining)>&& Receiver) override + virtual std::vector<std::function<void()>> GetLargeBuildBlob(const Oid& BuildId, + const IoHash& RawHash, + uint64_t ChunkSize, + std::function<void(uint64_t Offset, const IoBuffer& Chunk)>&& OnReceive, + std::function<void()>&& OnComplete) override { ZEN_TRACE_CPU("FileBuildStorage::GetLargeBuildBlob"); ZEN_UNUSED(BuildId); @@ -387,16 +387,18 @@ public: { struct WorkloadData { - std::atomic<uint64_t> BytesRemaining; - BasicFile BlobFile; - std::function<void(uint64_t Offset, const IoBuffer& Chunk, uint64_t BytesRemaining)> Receiver; + std::atomic<uint64_t> BytesRemaining; + BasicFile BlobFile; + std::function<void(uint64_t Offset, const IoBuffer& Chunk)> OnReceive; + std::function<void()> OnComplete; }; std::shared_ptr<WorkloadData> Workload(std::make_shared<WorkloadData>()); Workload->BlobFile.Open(BlockPath, BasicFile::Mode::kRead); const uint64_t BlobSize = Workload->BlobFile.FileSize(); - Workload->Receiver = std::move(Receiver); + Workload->OnReceive = std::move(OnReceive); + Workload->OnComplete = std::move(OnComplete); Workload->BytesRemaining = BlobSize; std::vector<std::function<void()>> WorkItems; @@ -410,8 +412,12 @@ public: IoBuffer PartPayload(Size); Workload->BlobFile.Read(PartPayload.GetMutableView().GetData(), Size, Offset); m_Stats.TotalBytesRead += PartPayload.GetSize(); + Workload->OnReceive(Offset, PartPayload); uint64_t ByteRemaning = Workload->BytesRemaining.fetch_sub(Size); - Workload->Receiver(Offset, PartPayload, ByteRemaning); + if (ByteRemaning == Size) + { + Workload->OnComplete(); + } SimulateLatency(Size, PartPayload.GetSize()); }); |