1#include "duckdb/storage/partial_block_manager.hpp"
2
3namespace duckdb {
4
5PartialBlockManager::PartialBlockManager(BlockManager &block_manager, CheckpointType checkpoint_type,
6 uint32_t max_partial_block_size, uint32_t max_use_count)
7 : block_manager(block_manager), checkpoint_type(checkpoint_type), max_partial_block_size(max_partial_block_size),
8 max_use_count(max_use_count) {
9}
10PartialBlockManager::~PartialBlockManager() {
11}
12//===--------------------------------------------------------------------===//
13// Partial Blocks
14//===--------------------------------------------------------------------===//
15PartialBlockAllocation PartialBlockManager::GetBlockAllocation(uint32_t segment_size) {
16 PartialBlockAllocation allocation;
17 allocation.block_manager = &block_manager;
18 allocation.allocation_size = segment_size;
19
20 // if the block is less than 80% full, we consider it a "partial block"
21 // which means we will try to fit it with other blocks
22 // check if there is a partial block available we can write to
23 if (segment_size <= max_partial_block_size && GetPartialBlock(segment_size, state&: allocation.partial_block)) {
24 //! there is! increase the reference count of this block
25 allocation.partial_block->state.block_use_count += 1;
26 allocation.state = allocation.partial_block->state;
27 if (checkpoint_type == CheckpointType::FULL_CHECKPOINT) {
28 block_manager.IncreaseBlockReferenceCount(block_id: allocation.state.block_id);
29 }
30 } else {
31 // full block: get a free block to write to
32 AllocateBlock(state&: allocation.state, segment_size);
33 }
34 return allocation;
35}
36
37bool PartialBlockManager::HasBlockAllocation(uint32_t segment_size) {
38 return segment_size <= max_partial_block_size &&
39 partially_filled_blocks.lower_bound(x: segment_size) != partially_filled_blocks.end();
40}
41
42void PartialBlockManager::AllocateBlock(PartialBlockState &state, uint32_t segment_size) {
43 D_ASSERT(segment_size <= Storage::BLOCK_SIZE);
44 if (checkpoint_type == CheckpointType::FULL_CHECKPOINT) {
45 state.block_id = block_manager.GetFreeBlockId();
46 } else {
47 state.block_id = INVALID_BLOCK;
48 }
49 state.block_size = Storage::BLOCK_SIZE;
50 state.offset_in_block = 0;
51 state.block_use_count = 1;
52}
53
54bool PartialBlockManager::GetPartialBlock(idx_t segment_size, unique_ptr<PartialBlock> &partial_block) {
55 auto entry = partially_filled_blocks.lower_bound(x: segment_size);
56 if (entry == partially_filled_blocks.end()) {
57 return false;
58 }
59 // found a partially filled block! fill in the info
60 partial_block = std::move(entry->second);
61 partially_filled_blocks.erase(position: entry);
62
63 D_ASSERT(partial_block->state.offset_in_block > 0);
64 D_ASSERT(ValueIsAligned(partial_block->state.offset_in_block));
65 return true;
66}
67
68void PartialBlockManager::RegisterPartialBlock(PartialBlockAllocation &&allocation) {
69 auto &state = allocation.partial_block->state;
70 if (state.block_use_count < max_use_count) {
71 auto unaligned_size = allocation.allocation_size + state.offset_in_block;
72 auto new_size = AlignValue(n: unaligned_size);
73 if (new_size != unaligned_size) {
74 // register the uninitialized region so we can correctly initialize it before writing to disk
75 allocation.partial_block->AddUninitializedRegion(start: unaligned_size, end: new_size);
76 }
77 state.offset_in_block = new_size;
78 auto new_space_left = state.block_size - new_size;
79 // check if the block is STILL partially filled after adding the segment_size
80 if (new_space_left >= Storage::BLOCK_SIZE - max_partial_block_size) {
81 // the block is still partially filled: add it to the partially_filled_blocks list
82 partially_filled_blocks.insert(x: make_pair(x&: new_space_left, y: std::move(allocation.partial_block)));
83 }
84 }
85 idx_t free_space = state.block_size - state.offset_in_block;
86 auto block_to_free = std::move(allocation.partial_block);
87 if (!block_to_free && partially_filled_blocks.size() > MAX_BLOCK_MAP_SIZE) {
88 // Free the page with the least space free.
89 auto itr = partially_filled_blocks.begin();
90 block_to_free = std::move(itr->second);
91 free_space = state.block_size - itr->first;
92 partially_filled_blocks.erase(position: itr);
93 }
94 // Flush any block that we're not going to reuse.
95 if (block_to_free) {
96 block_to_free->Flush(free_space_left: free_space);
97 AddWrittenBlock(block: block_to_free->state.block_id);
98 }
99}
100
101void PartialBlock::Merge(PartialBlock &other, idx_t offset, idx_t other_size) {
102 throw InternalException("PartialBlock::Merge not implemented for this block type");
103}
104
105void PartialBlockManager::Merge(PartialBlockManager &other) {
106 if (&other == this) {
107 throw InternalException("Cannot merge into itself");
108 }
109 // for each partially filled block in the other manager, check if we can merge it into an existing block in this
110 // manager
111 for (auto &e : other.partially_filled_blocks) {
112 if (!e.second) {
113 throw InternalException("Empty partially filled block found");
114 }
115 auto used_space = Storage::BLOCK_SIZE - e.first;
116 if (HasBlockAllocation(segment_size: used_space)) {
117 // we can merge this block into an existing block - merge them
118 // merge blocks
119 auto allocation = GetBlockAllocation(segment_size: used_space);
120 allocation.partial_block->Merge(other&: *e.second, offset: allocation.state.offset_in_block, other_size: used_space);
121
122 // re-register the partial block
123 allocation.state.offset_in_block += used_space;
124 RegisterPartialBlock(allocation: std::move(allocation));
125 } else {
126 // we cannot merge this block - append it directly to the current block manager
127 partially_filled_blocks.insert(x: make_pair(x: e.first, y: std::move(e.second)));
128 }
129 }
130 // copy over the written blocks
131 for (auto &block_id : other.written_blocks) {
132 AddWrittenBlock(block: block_id);
133 }
134 other.written_blocks.clear();
135 other.partially_filled_blocks.clear();
136}
137
138void PartialBlockManager::AddWrittenBlock(block_id_t block) {
139 auto entry = written_blocks.insert(x: block);
140 if (!entry.second) {
141 throw InternalException("Written block already exists");
142 }
143}
144
145void PartialBlockManager::ClearBlocks() {
146 for (auto &e : partially_filled_blocks) {
147 e.second->Clear();
148 }
149 partially_filled_blocks.clear();
150}
151
152void PartialBlockManager::FlushPartialBlocks() {
153 for (auto &e : partially_filled_blocks) {
154 e.second->Flush(free_space_left: e.first);
155 }
156 partially_filled_blocks.clear();
157}
158
159void PartialBlockManager::Rollback() {
160 ClearBlocks();
161 for (auto &block_id : written_blocks) {
162 block_manager.MarkBlockAsFree(block_id);
163 }
164}
165
166} // namespace duckdb
167