add multirange requests to blob store (#795)

* add multirange requests to blob store

3 files changed, 295 insertions, 21 deletions

diff --git a/src/zenhttp/packageformat.cpp b/src/zenhttp/packageformat.cpp
index 708238224..9a80d07c8 100644
--- a/src/zenhttp/packageformat.cpp
+++ b/src/zenhttp/packageformat.cpp
@@ -581,7 +581,7 @@ ParsePackageMessage(IoBuffer Payload, std::function<IoBuffer(const IoHash&, uint
 				ZEN_ASSERT(AttachmentBufferCopy.Size() == AttachmentSize);
 				AttachmentBufferCopy.GetMutableView().CopyFrom(AttachmentBuffer.GetView());
 
-				Attachments.emplace_back(SharedBuffer{AttachmentBufferCopy});
+				Attachments.emplace_back(CbAttachment(SharedBuffer{AttachmentBufferCopy}, Entry.AttachmentHash));
 			}
 			else
 			{
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 {
diff --git a/src/zenserver/storage/buildstore/httpbuildstore.cpp b/src/zenserver/storage/buildstore/httpbuildstore.cpp
index bf7afcc02..6ada085a5 100644
--- a/src/zenserver/storage/buildstore/httpbuildstore.cpp
+++ b/src/zenserver/storage/buildstore/httpbuildstore.cpp
@@ -71,7 +71,7 @@ HttpBuildStoreService::Initialize()
 	m_Router.RegisterRoute(
 		"{namespace}/{bucket}/{buildid}/blobs/{hash}",
 		[this](HttpRouterRequest& Req) { GetBlobRequest(Req); },
-		HttpVerb::kGet);
+		HttpVerb::kGet | HttpVerb::kPost);
 
 	m_Router.RegisterRoute(
 		"{namespace}/{bucket}/{buildid}/blobs/putBlobMetadata",
@@ -161,14 +161,49 @@ HttpBuildStoreService::GetBlobRequest(HttpRouterRequest& Req)
 										   HttpContentType::kText,
 										   fmt::format("Invalid blob hash '{}'", Hash));
 	}
-	zen::HttpRanges Ranges;
-	bool			HasRange = ServerRequest.TryGetRanges(Ranges);
-	if (Ranges.size() > 1)
+
+	std::vector<std::pair<uint64_t, uint64_t>> OffsetAndLengthPairs;
+	if (ServerRequest.RequestVerb() == HttpVerb::kPost)
 	{
-		// Only a single range is supported
-		return ServerRequest.WriteResponse(HttpResponseCode::BadRequest,
-										   HttpContentType::kText,
-										   "Multiple ranges in blob request is not supported");
+		CbObject RangePayload = ServerRequest.ReadPayloadObject();
+		if (RangePayload)
+		{
+			CbArrayView RangesArray = RangePayload["ranges"sv].AsArrayView();
+			OffsetAndLengthPairs.reserve(RangesArray.Num());
+			for (CbFieldView FieldView : RangesArray)
+			{
+				CbObjectView RangeView	 = FieldView.AsObjectView();
+				uint64_t	 RangeOffset = RangeView["offset"sv].AsUInt64();
+				uint64_t	 RangeLength = RangeView["length"sv].AsUInt64();
+				OffsetAndLengthPairs.push_back(std::make_pair(RangeOffset, RangeLength));
+			}
+		}
+		if (OffsetAndLengthPairs.empty())
+		{
+			m_BuildStoreStats.BadRequestCount++;
+			return ServerRequest.WriteResponse(HttpResponseCode::BadRequest,
+											   HttpContentType::kText,
+											   "Fetching blob without ranges must be done with the GET verb");
+		}
+	}
+	else
+	{
+		HttpRanges Ranges;
+		bool	   HasRange = ServerRequest.TryGetRanges(Ranges);
+		if (HasRange)
+		{
+			if (Ranges.size() > 1)
+			{
+				// Only a single http range is supported, we have limited support for http multirange responses
+				m_BuildStoreStats.BadRequestCount++;
+				return ServerRequest.WriteResponse(HttpResponseCode::BadRequest,
+												   HttpContentType::kText,
+												   fmt::format("Multiple ranges in blob request is only supported for {} accept type",
+															   ToString(HttpContentType::kCbPackage)));
+			}
+			const HttpRange& FirstRange = Ranges.front();
+			OffsetAndLengthPairs.push_back(std::make_pair<uint64_t, uint64_t>(FirstRange.Start, FirstRange.End - FirstRange.Start + 1));
+		}
 	}
 
 	m_BuildStoreStats.BlobReadCount++;
@@ -179,24 +214,65 @@ HttpBuildStoreService::GetBlobRequest(HttpRouterRequest& Req)
 										   HttpContentType::kText,
 										   fmt::format("Blob with hash '{}' could not be found", Hash));
 	}
-	//	ZEN_INFO("Fetched blob {}. Size: {}", BlobHash, Blob.GetSize());
 	m_BuildStoreStats.BlobHitCount++;
-	if (HasRange)
+
+	if (OffsetAndLengthPairs.empty())
 	{
-		const HttpRange& Range		 = Ranges.front();
-		const uint64_t	 BlobSize	 = Blob.GetSize();
-		const uint64_t	 MaxBlobSize = Range.Start < BlobSize ? BlobSize - Range.Start : 0;
-		const uint64_t	 RangeSize	 = Min(Range.End - Range.Start + 1, MaxBlobSize);
-		if (Range.Start + RangeSize > BlobSize)
+		return ServerRequest.WriteResponse(HttpResponseCode::OK, Blob.GetContentType(), Blob);
+	}
+
+	if (ServerRequest.AcceptContentType() == HttpContentType::kCbPackage)
+	{
+		const uint64_t BlobSize = Blob.GetSize();
+
+		CbPackage			  ResponsePackage;
+		std::vector<IoBuffer> RangeBuffers;
+		CbObjectWriter		  Writer;
+		Writer.BeginArray("ranges"sv);
+		for (const std::pair<uint64_t, uint64_t>& Range : OffsetAndLengthPairs)
 		{
-			return ServerRequest.WriteResponse(HttpResponseCode::NoContent);
+			const uint64_t MaxBlobSize = Range.first < BlobSize ? BlobSize - Range.first : 0;
+			const uint64_t RangeSize   = Min(Range.second, MaxBlobSize);
+			if (Range.first + RangeSize <= BlobSize)
+			{
+				RangeBuffers.push_back(IoBuffer(Blob, Range.first, RangeSize));
+				Writer.BeginObject();
+				{
+					Writer.AddInteger("offset"sv, Range.first);
+					Writer.AddInteger("length"sv, RangeSize);
+				}
+				Writer.EndObject();
+			}
 		}
-		Blob = IoBuffer(Blob, Range.Start, RangeSize);
-		return ServerRequest.WriteResponse(HttpResponseCode::OK, ZenContentType::kBinary, Blob);
+		Writer.EndArray();
+
+		CompositeBuffer Ranges(RangeBuffers);
+		CbAttachment	PayloadAttachment(std::move(Ranges), BlobHash);
+		Writer.AddAttachment("payload", PayloadAttachment);
+
+		CbObject HeaderObject = Writer.Save();
+
+		ResponsePackage.AddAttachment(PayloadAttachment);
+		ResponsePackage.SetObject(HeaderObject);
+
+		CompositeBuffer RpcResponseBuffer = FormatPackageMessageBuffer(ResponsePackage);
+		uint64_t		ResponseSize	  = RpcResponseBuffer.GetSize();
+		ZEN_UNUSED(ResponseSize);
+		return ServerRequest.WriteResponse(HttpResponseCode::OK, HttpContentType::kCbPackage, RpcResponseBuffer);
 	}
 	else
 	{
-		return ServerRequest.WriteResponse(HttpResponseCode::OK, Blob.GetContentType(), Blob);
+		ZEN_ASSERT(OffsetAndLengthPairs.size() == 1);
+		const std::pair<uint64_t, uint64_t>& OffsetAndLength = OffsetAndLengthPairs.front();
+		const uint64_t						 BlobSize		 = Blob.GetSize();
+		const uint64_t						 MaxBlobSize	 = OffsetAndLength.first < BlobSize ? BlobSize - OffsetAndLength.first : 0;
+		const uint64_t						 RangeSize		 = Min(OffsetAndLength.second, MaxBlobSize);
+		if (OffsetAndLength.first + RangeSize > BlobSize)
+		{
+			return ServerRequest.WriteResponse(HttpResponseCode::NoContent);
+		}
+		Blob = IoBuffer(Blob, OffsetAndLength.first, RangeSize);
+		return ServerRequest.WriteResponse(HttpResponseCode::OK, ZenContentType::kBinary, Blob);
 	}
 }