From efd016d84d0940bf616e3efef135532cbf5fedef Mon Sep 17 00:00:00 2001
From: Dan Engelbrecht <dan.engelbrecht@epicgames.com>
Date: Mon, 21 Oct 2024 15:40:13 +0200
Subject: bucket size queries (#203)

- Feature: Added options --bucketsize and --bucketsizes to zen cache-info to get data sizes in cache buckets and attachments
---
 src/zenserver/cache/httpstructuredcache.cpp | 115 ++++++++++++++++++++++++++++
 1 file changed, 115 insertions(+)

(limited to 'src/zenserver/cache/httpstructuredcache.cpp')

diff --git a/src/zenserver/cache/httpstructuredcache.cpp b/src/zenserver/cache/httpstructuredcache.cpp
index 109fb34f6..551b5a76d 100644
--- a/src/zenserver/cache/httpstructuredcache.cpp
+++ b/src/zenserver/cache/httpstructuredcache.cpp
@@ -23,6 +23,7 @@
 #include <zenutil/cache/cacherequests.h>
 #include <zenutil/cache/rpcrecording.h>
 #include <zenutil/packageformat.h>
+#include <zenutil/workerpools.h>
 
 #include "upstream/jupiter.h"
 #include "upstream/upstreamcache.h"
@@ -606,6 +607,81 @@ HttpStructuredCacheService::HandleCacheNamespaceRequest(HttpServerRequest& Reque
 
 				ResponseWriter.AddInteger("EntryCount", Info->DiskLayerInfo.EntryCount);
 
+				if (auto Buckets = HttpServerRequest::Decode(Request.GetQueryParams().GetValue("bucketsizes")); !Buckets.empty())
+				{
+					ResponseWriter.BeginObject("BucketSizes");
+
+					ResponseWriter.BeginArray("Buckets");
+
+					std::vector<std::string> BucketNames;
+					if (Buckets == "*")	 // Get all - empty FieldFilter equal getting all fields
+					{
+						BucketNames = Info.value().BucketNames;
+					}
+					else
+					{
+						ForEachStrTok(Buckets, ',', [&](std::string_view BucketName) {
+							BucketNames.push_back(std::string(BucketName));
+							return true;
+						});
+					}
+					WorkerThreadPool&	WorkerPool = GetMediumWorkerPool(EWorkloadType::Background);
+					std::vector<IoHash> AllAttachments;
+					for (const std::string& BucketName : BucketNames)
+					{
+						ResponseWriter.BeginObject();
+						ResponseWriter << "Name" << BucketName;
+						CacheContentStats ContentStats;
+						bool			  Success = m_CacheStore.GetContentStats(NamespaceName, BucketName, ContentStats);
+						if (Success)
+						{
+							size_t ValuesSize = 0;
+							for (const uint64_t Size : ContentStats.ValueSizes)
+							{
+								ValuesSize += Size;
+							}
+
+							std::sort(ContentStats.Attachments.begin(), ContentStats.Attachments.end());
+							auto NewEnd = std::unique(ContentStats.Attachments.begin(), ContentStats.Attachments.end());
+							ContentStats.Attachments.erase(NewEnd, ContentStats.Attachments.end());
+
+							ResponseWriter << "Count" << ContentStats.ValueSizes.size();
+							ResponseWriter << "StructuredCount" << ContentStats.StructuredValuesCount;
+							ResponseWriter << "StandaloneCount" << ContentStats.StandaloneValuesCount;
+							ResponseWriter << "Size" << ValuesSize;
+							ResponseWriter << "AttachmentCount" << ContentStats.Attachments.size();
+
+							AllAttachments.insert(AllAttachments.end(), ContentStats.Attachments.begin(), ContentStats.Attachments.end());
+						}
+						ResponseWriter.EndObject();
+					}
+
+					ResponseWriter.EndArray();
+
+					ResponseWriter.BeginObject("Attachments");
+					std::sort(AllAttachments.begin(), AllAttachments.end());
+					auto NewEnd = std::unique(AllAttachments.begin(), AllAttachments.end());
+					AllAttachments.erase(NewEnd, AllAttachments.end());
+
+					uint64_t AttachmentsSize = 0;
+
+					m_CidStore.IterateChunks(
+						AllAttachments,
+						[&](size_t Index, const IoBuffer& Payload) {
+							ZEN_UNUSED(Index);
+							AttachmentsSize += Payload.GetSize();
+							return true;
+						},
+						&WorkerPool);
+
+					ResponseWriter << "Count" << AllAttachments.size();
+					ResponseWriter << "Size" << AttachmentsSize;
+
+					ResponseWriter.EndObject();
+
+					ResponseWriter.EndObject();
+				}
+
 				return Request.WriteResponse(HttpResponseCode::OK, ResponseWriter.Save());
 			}
 			break;
@@ -656,6 +732,45 @@ HttpStructuredCacheService::HandleCacheBucketRequest(HttpServerRequest& Request,
 
 				ResponseWriter.AddInteger("DiskEntryCount", Info->DiskLayerInfo.EntryCount);
 
+				if (auto GetBucketSize = Request.GetQueryParams().GetValue("bucketsize"); GetBucketSize == "true")
+				{
+					CacheContentStats ContentStats;
+					bool			  Success = m_CacheStore.GetContentStats(NamespaceName, BucketName, ContentStats);
+					if (Success)
+					{
+						size_t ValuesSize = 0;
+						for (const uint64_t Size : ContentStats.ValueSizes)
+						{
+							ValuesSize += Size;
+						}
+
+						std::sort(ContentStats.Attachments.begin(), ContentStats.Attachments.end());
+						auto NewEnd = std::unique(ContentStats.Attachments.begin(), ContentStats.Attachments.end());
+						ContentStats.Attachments.erase(NewEnd, ContentStats.Attachments.end());
+
+						ResponseWriter << "Count" << ContentStats.ValueSizes.size();
+						ResponseWriter << "StructuredCount" << ContentStats.StructuredValuesCount;
+						ResponseWriter << "StandaloneCount" << ContentStats.StandaloneValuesCount;
+						ResponseWriter << "Size" << ValuesSize;
+						ResponseWriter << "AttachmentCount" << ContentStats.Attachments.size();
+
+						uint64_t AttachmentsSize = 0;
+
+						WorkerThreadPool& WorkerPool = GetMediumWorkerPool(EWorkloadType::Background);
+
+						m_CidStore.IterateChunks(
+							ContentStats.Attachments,
+							[&](size_t Index, const IoBuffer& Payload) {
+								ZEN_UNUSED(Index);
+								AttachmentsSize += Payload.GetSize();
+								return true;
+							},
+							&WorkerPool);
+
+						ResponseWriter << "AttachmentsSize" << AttachmentsSize;
+					}
+				}
+
 				return Request.WriteResponse(HttpResponseCode::OK, ResponseWriter.Save());
 			}
 			break;
-- 
cgit v1.2.3