add multirange requests to blob store (#795)

* add multirange requests to blob store

1 files changed, 199 insertions, 1 deletions

diff --git a/src/zenserver-test/buildstore-tests.cpp b/src/zenserver-test/buildstore-tests.cpp
index 02b308485..ef48b2362 100644
--- a/src/zenserver-test/buildstore-tests.cpp
+++ b/src/zenserver-test/buildstore-tests.cpp
@@ -36,7 +36,8 @@ TEST_CASE("buildstore.blobs")
 	std::string_view Bucket	   = "bkt"sv;
 	Oid				 BuildId   = Oid::NewOid();
 
-	std::vector<IoHash> CompressedBlobsHashes;
+	std::vector<IoHash>	  CompressedBlobsHashes;
+	std::vector<uint64_t> CompressedBlobsSizes;
 	{
 		ZenServerInstance Instance(TestEnv);
 
@@ -51,6 +52,7 @@ TEST_CASE("buildstore.blobs")
 			IoBuffer		 Blob			= CreateSemiRandomBlob(4711 + I * 7);
 			CompressedBuffer CompressedBlob = CompressedBuffer::Compress(SharedBuffer(std::move(Blob)));
 			CompressedBlobsHashes.push_back(CompressedBlob.DecodeRawHash());
+			CompressedBlobsSizes.push_back(CompressedBlob.GetCompressedSize());
 			IoBuffer Payload = std::move(CompressedBlob).GetCompressed().Flatten().AsIoBuffer();
 			Payload.SetContentType(ZenContentType::kCompressedBinary);
 
@@ -107,6 +109,7 @@ TEST_CASE("buildstore.blobs")
 			IoBuffer		 Blob			= CreateSemiRandomBlob(5713 + I * 7);
 			CompressedBuffer CompressedBlob = CompressedBuffer::Compress(SharedBuffer(std::move(Blob)));
 			CompressedBlobsHashes.push_back(CompressedBlob.DecodeRawHash());
+			CompressedBlobsSizes.push_back(CompressedBlob.GetCompressedSize());
 			IoBuffer Payload = std::move(CompressedBlob).GetCompressed().Flatten().AsIoBuffer();
 			Payload.SetContentType(ZenContentType::kCompressedBinary);
 
@@ -141,6 +144,201 @@ TEST_CASE("buildstore.blobs")
 			CHECK(IoHash::HashBuffer(Decompressed) == RawHash);
 		}
 	}
+
+	{
+		// Single-range Get
+
+		ZenServerInstance Instance(TestEnv);
+
+		const uint16_t PortNumber =
+			Instance.SpawnServerAndWaitUntilReady(fmt::format("--buildstore-enabled --system-dir {}", SystemRootPath));
+		CHECK(PortNumber != 0);
+
+		HttpClient Client(Instance.GetBaseUri() + "/builds/");
+
+		{
+			const IoHash& RawHash  = CompressedBlobsHashes.front();
+			uint64_t	  BlobSize = CompressedBlobsSizes.front();
+
+			std::vector<std::pair<uint64_t, uint64_t>> Ranges = {{BlobSize / 16 * 1, BlobSize / 2}};
+
+			uint64_t RangeSizeSum = Ranges.front().second;
+
+			HttpClient::KeyValueMap Headers;
+
+			Headers.Entries.insert(
+				{"Range", fmt::format("bytes={}-{}", Ranges.front().first, Ranges.front().first + Ranges.front().second - 1)});
+
+			HttpClient::Response Result = Client.Get(fmt::format("{}/{}/{}/blobs/{}", Namespace, Bucket, BuildId, RawHash), Headers);
+			REQUIRE(Result);
+			IoBuffer Payload = Result.ResponsePayload;
+			CHECK_EQ(RangeSizeSum, Payload.GetSize());
+
+			HttpClient::Response FullBlobResult = Client.Get(fmt::format("{}/{}/{}/blobs/{}", Namespace, Bucket, BuildId, RawHash),
+															 HttpClient::Accept(ZenContentType::kCompressedBinary));
+			REQUIRE(FullBlobResult);
+			MemoryView ActualRange = FullBlobResult.ResponsePayload.GetView().Mid(Ranges.front().first, Ranges.front().second);
+			MemoryView RangeView   = Payload.GetView();
+			CHECK(ActualRange.EqualBytes(RangeView));
+		}
+	}
+
+	{
+		// Single-range Post
+
+		ZenServerInstance Instance(TestEnv);
+
+		const uint16_t PortNumber =
+			Instance.SpawnServerAndWaitUntilReady(fmt::format("--buildstore-enabled --system-dir {}", SystemRootPath));
+		CHECK(PortNumber != 0);
+
+		HttpClient Client(Instance.GetBaseUri() + "/builds/");
+
+		{
+			uint64_t RangeSizeSum = 0;
+
+			const IoHash& RawHash  = CompressedBlobsHashes.front();
+			uint64_t	  BlobSize = CompressedBlobsSizes.front();
+
+			std::vector<std::pair<uint64_t, uint64_t>> Ranges = {{BlobSize / 16 * 1, BlobSize / 2}};
+
+			CbObjectWriter Writer;
+			Writer.BeginArray("ranges"sv);
+			{
+				for (const std::pair<uint64_t, uint64_t>& Range : Ranges)
+				{
+					Writer.BeginObject();
+					{
+						Writer.AddInteger("offset"sv, Range.first);
+						Writer.AddInteger("length"sv, Range.second);
+						RangeSizeSum += Range.second;
+					}
+					Writer.EndObject();
+				}
+			}
+			Writer.EndArray();	// ranges
+
+			HttpClient::Response Result = Client.Post(fmt::format("{}/{}/{}/blobs/{}", Namespace, Bucket, BuildId, RawHash),
+													  Writer.Save(),
+													  HttpClient::Accept(ZenContentType::kCbPackage));
+			REQUIRE(Result);
+			IoBuffer Payload = Result.ResponsePayload;
+			REQUIRE(Payload.GetContentType() == ZenContentType::kCbPackage);
+
+			CbPackage	 ResponsePackage = ParsePackageMessage(Payload);
+			CbObjectView ResponseObject	 = ResponsePackage.GetObject();
+
+			CbArrayView RangeArray = ResponseObject["ranges"sv].AsArrayView();
+			CHECK_EQ(RangeArray.Num(), Ranges.size());
+			size_t RangeOffset = 0;
+			for (CbFieldView View : RangeArray)
+			{
+				CbObjectView Range = View.AsObjectView();
+				CHECK_EQ(Range["offset"sv].AsUInt64(), Ranges[RangeOffset].first);
+				CHECK_EQ(Range["length"sv].AsUInt64(), Ranges[RangeOffset].second);
+				RangeOffset++;
+			}
+
+			const CbAttachment* DataAttachment = ResponsePackage.FindAttachment(RawHash);
+			REQUIRE(DataAttachment);
+			SharedBuffer PayloadRanges = DataAttachment->AsBinary();
+			CHECK_EQ(RangeSizeSum, PayloadRanges.GetSize());
+
+			HttpClient::Response FullBlobResult = Client.Get(fmt::format("{}/{}/{}/blobs/{}", Namespace, Bucket, BuildId, RawHash),
+															 HttpClient::Accept(ZenContentType::kCompressedBinary));
+			REQUIRE(FullBlobResult);
+
+			uint64_t Offset = 0;
+			for (const std::pair<uint64_t, uint64_t>& Range : Ranges)
+			{
+				MemoryView ActualRange = FullBlobResult.ResponsePayload.GetView().Mid(Range.first, Range.second);
+				MemoryView RangeView   = PayloadRanges.GetView().Mid(Offset, Range.second);
+				CHECK(ActualRange.EqualBytes(RangeView));
+				Offset += Range.second;
+			}
+		}
+	}
+
+	{
+		// Multi-range
+
+		ZenServerInstance Instance(TestEnv);
+
+		const uint16_t PortNumber =
+			Instance.SpawnServerAndWaitUntilReady(fmt::format("--buildstore-enabled --system-dir {}", SystemRootPath));
+		CHECK(PortNumber != 0);
+
+		HttpClient Client(Instance.GetBaseUri() + "/builds/");
+
+		{
+			uint64_t RangeSizeSum = 0;
+
+			const IoHash& RawHash  = CompressedBlobsHashes.front();
+			uint64_t	  BlobSize = CompressedBlobsSizes.front();
+
+			std::vector<std::pair<uint64_t, uint64_t>> Ranges = {
+				{BlobSize / 16 * 1, BlobSize / 20},
+				{BlobSize / 16 * 3, BlobSize / 32},
+				{BlobSize / 16 * 5, BlobSize / 16},
+				{BlobSize - BlobSize / 16, BlobSize / 16 - 1},
+			};
+
+			CbObjectWriter Writer;
+			Writer.BeginArray("ranges"sv);
+			{
+				for (const std::pair<uint64_t, uint64_t>& Range : Ranges)
+				{
+					Writer.BeginObject();
+					{
+						Writer.AddInteger("offset"sv, Range.first);
+						Writer.AddInteger("length"sv, Range.second);
+						RangeSizeSum += Range.second;
+					}
+					Writer.EndObject();
+				}
+			}
+			Writer.EndArray();	// ranges
+
+			HttpClient::Response Result = Client.Post(fmt::format("{}/{}/{}/blobs/{}", Namespace, Bucket, BuildId, RawHash),
+													  Writer.Save(),
+													  HttpClient::Accept(ZenContentType::kCbPackage));
+			REQUIRE(Result);
+			IoBuffer Payload = Result.ResponsePayload;
+			REQUIRE(Payload.GetContentType() == ZenContentType::kCbPackage);
+
+			CbPackage	 ResponsePackage = ParsePackageMessage(Payload);
+			CbObjectView ResponseObject	 = ResponsePackage.GetObject();
+
+			CbArrayView RangeArray = ResponseObject["ranges"sv].AsArrayView();
+			CHECK_EQ(RangeArray.Num(), Ranges.size());
+			size_t RangeOffset = 0;
+			for (CbFieldView View : RangeArray)
+			{
+				CbObjectView Range = View.AsObjectView();
+				CHECK_EQ(Range["offset"sv].AsUInt64(), Ranges[RangeOffset].first);
+				CHECK_EQ(Range["length"sv].AsUInt64(), Ranges[RangeOffset].second);
+				RangeOffset++;
+			}
+
+			const CbAttachment* DataAttachment = ResponsePackage.FindAttachment(RawHash);
+			REQUIRE(DataAttachment);
+			SharedBuffer PayloadRanges = DataAttachment->AsBinary();
+			CHECK_EQ(RangeSizeSum, PayloadRanges.GetSize());
+
+			HttpClient::Response FullBlobResult = Client.Get(fmt::format("{}/{}/{}/blobs/{}", Namespace, Bucket, BuildId, RawHash),
+															 HttpClient::Accept(ZenContentType::kCompressedBinary));
+			REQUIRE(FullBlobResult);
+
+			uint64_t Offset = 0;
+			for (const std::pair<uint64_t, uint64_t>& Range : Ranges)
+			{
+				MemoryView ActualRange = FullBlobResult.ResponsePayload.GetView().Mid(Range.first, Range.second);
+				MemoryView RangeView   = PayloadRanges.GetView().Mid(Offset, Range.second);
+				CHECK(ActualRange.EqualBytes(RangeView));
+				Offset += Range.second;
+			}
+		}
+	}
 }
 
 namespace {