Unload unused data from DataStorage

cmd/prometheus/main.go

@@ -771,6 +771,7 @@ func main() {
 		// x3 ScrapeInterval timeout for write block
 		time.Duration(cfgFile.GlobalConfig.ScrapeInterval*3),
 		cfg.WalMaxSamplesPerSegment,
+		appender.UnloadDataStorage,
 	)
 	if err != nil {
 		level.Error(logger).Log("msg", "failed to create a receiver", "err", err)
@@ -2105,6 +2106,10 @@ func readPromPPFeatures(logger log.Logger) {
 		case "disable_commits_on_remote_write":
 			rwprocessor.AlwaysCommit = false
 			pphandler.OTLPAlwaysCommit = false
+
+		case "unload_data_storage":
+			appender.UnloadDataStorage = true
+			_ = level.Info(logger).Log("msg", "[FEATURE] Data storage unloading is enabled.")
 		}
 	}
 }

cmd/prompptool/walpp.go

@@ -84,7 +84,7 @@ func (cmd *cmdWALPPToBlock) Do(
 	level.Debug(logger).Log("msg", "catalog records", "len", len(headRecords))
 
 	var inputRelabelerConfig []*config.InputRelabelerConfig
-	bw := block.NewBlockWriter(workingDir, block.DefaultChunkSegmentSize, time.Duration(cmd.blockDuration), registerer)
+	bw := block.NewWriter(workingDir, block.DefaultChunkSegmentSize, time.Duration(cmd.blockDuration), registerer)
 	for _, headRecord := range headRecords {
 		if err := ctx.Err(); err != nil {
 			return err
@@ -99,6 +99,7 @@ func (cmd *cmdWALPPToBlock) Do(
 			0,
 			head.NoOpLastAppendedSegmentIDSetter{},
 			registerer,
+			time.Duration(0),
 		)
 		if err != nil {
 			level.Error(logger).Log(
@@ -116,10 +117,8 @@ func (cmd *cmdWALPPToBlock) Do(
 		tBlockWrite := h.CreateTask(
 			relabeler.BlockWrite,
 			func(shard relabeler.Shard) error {
-				shard.LSSLock()
-				defer shard.LSSUnlock()
-
-				return bw.Write(relabeler.NewBlock(shard.LSS().Raw(), shard.DataStorage().Raw()))
+				_, err := bw.Write(shard)
+				return err
 			},
 			relabeler.ForLSSTask,
 		)

pp-pkg/receiver/receiver.go

@@ -125,6 +125,7 @@ func NewReceiver(
 	headRetentionTimeout time.Duration,
 	writeTimeout time.Duration,
 	maxSegmentSize uint32,
+	unloadDataStorage bool,
 ) (*Receiver, error) {
 	if logger == nil {
 		logger = log.NewNopLogger()
@@ -170,25 +171,31 @@ func NewReceiver(
 		return nil, err
 	}
 
+	var unloadDataStorageInterval time.Duration
+	if unloadDataStorage {
+		unloadDataStorageInterval = appender.DefaultMergeDuration
+	}
+
 	headManager, err := headmanager.New(
 		dataDir,
 		clock,
 		headConfigStorage,
 		headCatalog,
 		maxSegmentSize,
 		registerer,
+		unloadDataStorageInterval,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create head manager: %w", err)
 	}
 
-	activeHead, rotatedHeads, err := headManager.Restore(rotationInfo.BlockDuration)
+	activeHead, rotatedHeads, err := headManager.Restore(rotationInfo.BlockDuration, unloadDataStorageInterval)
 	if err != nil {
 		return nil, fmt.Errorf("failed to restore heads: %w", err)
 	}
 	readyNotifier.NotifyReady()
 	queryableStorage := appender.NewQueryableStorageWithWriteNotifier(
-		block.NewBlockWriter(dataDir, block.DefaultChunkSegmentSize, rotationInfo.BlockDuration, registerer),
+		block.NewWriter(dataDir, block.DefaultChunkSegmentSize, rotationInfo.BlockDuration, registerer),
 		registerer,
 		querier.NewMetrics(registerer, querier.QueryableStorageSource),
 		triggerNotifier,
@@ -228,6 +235,7 @@ func NewReceiver(
 			relabeler.NewRotateTimerWithSeed(clock, rotationInfo.BlockDuration, rotationInfo.Seed),
 			appender.NewConstantIntervalTimer(clock, commitInterval),
 			appender.NewConstantIntervalTimer(clock, appender.DefaultMergeDuration),
+			unloadDataStorage,
 			registerer,
 		),
 

pp/bare_bones/bitset.h

@@ -14,6 +14,7 @@
 #include <numeric>
 
 #include "bit.h"
+#include "concepts.h"
 #include "memory.h"
 #include "streams.h"
 #include "type_traits.h"
@@ -38,7 +39,7 @@ class Bitset {
     if (__builtin_expect(size > std::numeric_limits<uint32_t>::max(), false))
       std::abort();
 
-    const uint64_t size_in_uint64_elements = Bit::to_ceil_units<uint64_t>(size);
+    const auto size_in_uint64_elements = Bit::to_ceil_units<uint64_t>(size);
 
     if (size_in_uint64_elements <= data_.size()) {
       return;
@@ -186,21 +187,26 @@ class Bitset {
   }
 
   [[nodiscard]] PROMPP_ALWAYS_INLINE uint32_t get_write_size() const noexcept {
-    const uint32_t data_size_in_bytes = Bit::to_ceil_units<uint64_t>(size()) * sizeof(uint64_t);
+    const uint32_t data_size_in_bytes = memory_size_in_bytes();
     return sizeof(data_size_in_bytes) + data_size_in_bytes;
   }
 
   template <OutputStream S>
   PROMPP_ALWAYS_INLINE void write_to(S& stream) const noexcept {
     const uint32_t data_size_in_bits = size();
-    const uint32_t data_size_in_bytes = Bit::to_ceil_units<uint64_t>(data_size_in_bits) * sizeof(uint64_t);
+    const uint32_t data_size_in_bytes = memory_size_in_bytes();
+
+    if constexpr (BareBones::concepts::has_reserve<S>) {
+      stream.reserve(sizeof(data_size_in_bits) + data_size_in_bytes);
+    }
+
     stream.write(reinterpret_cast<const char*>(&data_size_in_bits), sizeof(data_size_in_bits));
     stream.write(reinterpret_cast<const char*>(data_.begin()), data_size_in_bytes);
   }
 
   static PROMPP_ALWAYS_INLINE Iterator create_read_iterator(std::span<const uint8_t>& buffer) noexcept {
     if (buffer.size() < sizeof(uint32_t)) [[unlikely]] {
-      return Iterator{};
+      return Iterator{nullptr, 0, 0};
     }
 
     uint32_t bit_count = 0;
@@ -210,17 +216,40 @@ class Bitset {
     const uint32_t uint64_count = BareBones::Bit::to_ceil_units<uint64_t>(bit_count);
     const uint32_t byte_count = uint64_count * sizeof(uint64_t);
     if (buffer.size() < byte_count) [[unlikely]] {
-      return Iterator{};
+      return Iterator{nullptr, 0, 0};
     }
 
     const std::span bit_data(reinterpret_cast<const uint64_t*>(buffer.data()), uint64_count);
     buffer = buffer.subspan(byte_count);
 
-    return Iterator(bit_data.data(), bit_count);
+    return Iterator(bit_data.data(), bit_count, 0);
+  }
+
+  [[nodiscard]] bool read_from(std::istream& stream) {
+    uint32_t bit_count{};
+    stream.read(reinterpret_cast<char*>(&bit_count), sizeof(uint32_t));
+    if (stream.gcount() != sizeof(uint32_t)) [[unlikely]] {
+      return false;
+    }
+
+    if (bit_count == 0) {
+      return true;
+    }
+
+    resize(bit_count);
+    const auto size_in_bytes = static_cast<std::streamsize>(memory_size_in_bytes(bit_count));
+    stream.read(reinterpret_cast<char*>(data_.begin()), size_in_bytes);
+
+    return stream.gcount() == size_in_bytes;
   }
 
  private:
   PROMPP_ALWAYS_INLINE void set_size(uint32_t new_size) noexcept { data_.control_block().items_count = new_size; }
+
+  [[nodiscard]] static PROMPP_ALWAYS_INLINE size_t memory_size_in_bytes(uint32_t bytes) noexcept {
+    return Bit::to_ceil_units<uint64_t>(bytes) * sizeof(uint64_t);
+  }
+  [[nodiscard]] PROMPP_ALWAYS_INLINE size_t memory_size_in_bytes() const noexcept { return memory_size_in_bytes(size()); }
 };
 
 template <>

pp/bare_bones/iterator.h

@@ -1,6 +1,11 @@
 #pragma once
 
+#include <algorithm>
+#include <cassert>
 #include <cstddef>
+#include <cstdint>
+
+#include "preprocess.h"
 
 namespace BareBones::iterator {
 
@@ -19,4 +24,45 @@ class OperationIterator {
   Operation* operation_;
 };
 
+template <class Iterator, class IteratorSentinel>
+class BatchIterator {
+ public:
+  using difference_type = typename Iterator::difference_type;
+  using iterator_category = std::input_iterator_tag;
+  using value_type = typename Iterator::value_type;
+  using pointer = value_type*;
+  using reference = value_type&;
+
+  PROMPP_ALWAYS_INLINE BatchIterator(Iterator iterator, uint32_t batch_size) noexcept : iterator_(std::move(iterator)), batch_size_(batch_size) {
+    assert(batch_size > 0);
+  }
+
+  [[nodiscard]] uint32_t batch_size() const noexcept { return batch_size_; }
+  void next_batch() { processed_ = 0; }
+
+  PROMPP_ALWAYS_INLINE BatchIterator& operator++() noexcept {
+    ++processed_;
+    ++iterator_;
+    return *this;
+  }
+
+  PROMPP_ALWAYS_INLINE BatchIterator operator++(int) noexcept {
+    const auto result = *this;
+    ++*this;
+    return result;
+  }
+
+  auto operator*() const noexcept { return *iterator_; }
+  auto operator*() noexcept { return *iterator_; }
+  auto operator->() const noexcept { return iterator_.operator->(); }
+  auto operator->() noexcept { return iterator_.operator->(); }
+
+  PROMPP_ALWAYS_INLINE bool operator==(const IteratorSentinel& end) const noexcept { return processed_ == batch_size_ || iterator_ == end; }
+
+ private:
+  Iterator iterator_;
+  uint32_t batch_size_;
+  uint32_t processed_{};
+};
+
 };  // namespace BareBones::iterator

pp/bare_bones/tests/bitset_tests.cpp

@@ -1,5 +1,6 @@
 #include <algorithm>
 #include <random>
+#include <spanstream>
 #include <string>
 #include <vector>
 
@@ -9,6 +10,8 @@
 
 namespace {
 
+using std::operator""sv;
+
 class BitsetFixture : public testing::Test {
  protected:
   BareBones::Bitset bs_;
@@ -358,4 +361,69 @@ TEST_F(BitsetCreateIteratorValidFixture, CreateReadIteratorValid) {
   EXPECT_EQ(buffer.size(), 0);
 }
 
+class BitsetReadFromFixture : public testing::Test {
+ protected:
+  BareBones::Bitset bs_;
+};
+
+TEST_F(BitsetReadFromFixture, ReadSizeError) {
+  // Arrange
+  std::ispanstream stream{""};
+
+  // Act
+  const auto result = bs_.read_from(stream);
+
+  // Assert
+  EXPECT_FALSE(result);
+}
+
+TEST_F(BitsetReadFromFixture, ReadBytesError) {
+  // Arrange
+  std::ispanstream stream{"\x01\x00\x00\x00"sv};
+
+  // Act
+  const auto result = bs_.read_from(stream);
+
+  // Assert
+  EXPECT_FALSE(result);
+}
+
+TEST_F(BitsetReadFromFixture, ReadEmptyBitset) {
+  // Arrange
+  std::ispanstream stream{"\x00\x00\x00\x00"sv};
+
+  // Act
+  const auto result = bs_.read_from(stream);
+
+  // Assert
+  EXPECT_TRUE(result);
+}
+
+TEST_F(BitsetReadFromFixture, ReadSuccess) {
+  // Arrange
+  std::ispanstream stream{"\x01\x00\x00\x00\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF"sv};
+
+  // Act
+  const auto result = bs_.read_from(stream);
+
+  // Assert
+  EXPECT_TRUE(result);
+  EXPECT_EQ(1U, bs_.size());
+  EXPECT_TRUE(bs_.is_set(0));
+}
+
+TEST_F(BitsetReadFromFixture, OverwriteBitsetAfterRead) {
+  // Arrange
+  bs_.set({1, 2, 3, 4, 5});
+  std::ispanstream stream{"\x01\x00\x00\x00\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF"sv};
+
+  // Act
+  const auto result = bs_.read_from(stream);
+
+  // Assert
+  EXPECT_TRUE(result);
+  EXPECT_EQ(1U, bs_.size());
+  EXPECT_TRUE(bs_.is_set(0));
+}
+
 }  // namespace