1//===----------------------------------------------------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9/// \file
10/// This file implements OnDiskGraphDB, an on-disk CAS nodes database,
11/// independent of a particular hashing algorithm. It only needs to be
12/// configured for the hash size and controls the schema of the storage.
13///
14/// OnDiskGraphDB defines:
15///
16/// - How the data is stored inside database, either as a standalone file, or
17/// allocated inside a datapool.
18/// - How references to other objects inside the same database is stored. They
19/// are stored as internal references, instead of full hash value to save
20/// space.
21/// - How to chain databases together and import objects from upstream
22/// databases.
23///
24/// Here's a top-level description of the current layout:
25///
26/// - db/index.<version>: a file for the "index" table, named by \a
27/// IndexTableName and managed by \a TrieRawHashMap. The contents are 8B
28/// that are accessed atomically, describing the object kind and where/how
29/// it's stored (including an optional file offset). See \a TrieRecord for
30/// more details.
31/// - db/data.<version>: a file for the "data" table, named by \a
32/// DataPoolTableName and managed by \a DataStore. New objects within
33/// TrieRecord::MaxEmbeddedSize are inserted here as \a
34/// TrieRecord::StorageKind::DataPool.
35/// - db/obj.<offset>.<version>: a file storing an object outside the main
36/// "data" table, named by its offset into the "index" table, with the
37/// format of \a TrieRecord::StorageKind::Standalone.
38/// - db/leaf.<offset>.<version>: a file storing a leaf node outside the
39/// main "data" table, named by its offset into the "index" table, with
40/// the format of \a TrieRecord::StorageKind::StandaloneLeaf.
41/// - db/leaf+0.<offset>.<version>: a file storing a null-terminated leaf object
42/// outside the main "data" table, named by its offset into the "index" table,
43/// with the format of \a TrieRecord::StorageKind::StandaloneLeaf0.
44//
45//===----------------------------------------------------------------------===//
46
47#include "llvm/CAS/OnDiskGraphDB.h"
48#include "OnDiskCommon.h"
49#include "llvm/ADT/DenseMap.h"
50#include "llvm/ADT/ScopeExit.h"
51#include "llvm/ADT/StringExtras.h"
52#include "llvm/CAS/OnDiskCASLogger.h"
53#include "llvm/CAS/OnDiskDataAllocator.h"
54#include "llvm/CAS/OnDiskTrieRawHashMap.h"
55#include "llvm/Support/Alignment.h"
56#include "llvm/Support/Compiler.h"
57#include "llvm/Support/Errc.h"
58#include "llvm/Support/Error.h"
59#include "llvm/Support/ErrorHandling.h"
60#include "llvm/Support/FileSystem.h"
61#include "llvm/Support/IOSandbox.h"
62#include "llvm/Support/MemoryBuffer.h"
63#include "llvm/Support/Path.h"
64#include "llvm/Support/Process.h"
65#include <atomic>
66#include <mutex>
67#include <optional>
68#include <variant>
69
70#define DEBUG_TYPE "on-disk-cas"
71
72using namespace llvm;
73using namespace llvm::cas;
74using namespace llvm::cas::ondisk;
75
76static constexpr StringLiteral IndexTableName = "llvm.cas.index";
77static constexpr StringLiteral DataPoolTableName = "llvm.cas.data";
78
79static constexpr StringLiteral IndexFilePrefix = "index.";
80static constexpr StringLiteral DataPoolFilePrefix = "data.";
81
82static constexpr StringLiteral FilePrefixObject = "obj.";
83static constexpr StringLiteral FilePrefixLeaf = "leaf.";
84static constexpr StringLiteral FilePrefixLeaf0 = "leaf+0.";
85
86static Error createCorruptObjectError(Expected<ArrayRef<uint8_t>> ID) {
87 if (!ID)
88 return ID.takeError();
89
90 return createStringError(EC: llvm::errc::invalid_argument,
91 S: "corrupt object '" + toHex(Input: *ID) + "'");
92}
93
94namespace {
95
96/// Trie record data: 8 bytes, atomic<uint64_t>
97/// - 1-byte: StorageKind
98/// - 7-bytes: DataStoreOffset (offset into referenced file)
99class TrieRecord {
100public:
101 enum class StorageKind : uint8_t {
102 /// Unknown object.
103 Unknown = 0,
104
105 /// data.vX: main pool, full DataStore record.
106 DataPool = 1,
107
108 /// obj.<TrieRecordOffset>.vX: standalone, with a full DataStore record.
109 Standalone = 10,
110
111 /// leaf.<TrieRecordOffset>.vX: standalone, just the data. File contents
112 /// exactly the data content and file size matches the data size. No refs.
113 StandaloneLeaf = 11,
114
115 /// leaf+0.<TrieRecordOffset>.vX: standalone, just the data plus an
116 /// extra null character ('\0'). File size is 1 bigger than the data size.
117 /// No refs.
118 StandaloneLeaf0 = 12,
119 };
120
121 static StringRef getStandaloneFilePrefix(StorageKind SK) {
122 switch (SK) {
123 default:
124 llvm_unreachable("Expected standalone storage kind");
125 case TrieRecord::StorageKind::Standalone:
126 return FilePrefixObject;
127 case TrieRecord::StorageKind::StandaloneLeaf:
128 return FilePrefixLeaf;
129 case TrieRecord::StorageKind::StandaloneLeaf0:
130 return FilePrefixLeaf0;
131 }
132 }
133
134 enum Limits : int64_t {
135 /// Saves files bigger than 64KB standalone instead of embedding them.
136 MaxEmbeddedSize = 64LL * 1024LL - 1,
137 };
138
139 struct Data {
140 StorageKind SK = StorageKind::Unknown;
141 FileOffset Offset;
142 };
143
144 /// Pack StorageKind and Offset from Data into 8 byte TrieRecord.
145 static uint64_t pack(Data D) {
146 assert(D.Offset.get() < (int64_t)(1ULL << 56));
147 uint64_t Packed = uint64_t(D.SK) << 56 | D.Offset.get();
148 assert(D.SK != StorageKind::Unknown || Packed == 0);
149#ifndef NDEBUG
150 Data RoundTrip = unpack(Packed);
151 assert(D.SK == RoundTrip.SK);
152 assert(D.Offset.get() == RoundTrip.Offset.get());
153#endif
154 return Packed;
155 }
156
157 // Unpack TrieRecord into Data.
158 static Data unpack(uint64_t Packed) {
159 Data D;
160 if (!Packed)
161 return D;
162 D.SK = (StorageKind)(Packed >> 56);
163 D.Offset = FileOffset(Packed & (UINT64_MAX >> 8));
164 return D;
165 }
166
167 TrieRecord() : Storage(0) {}
168
169 Data load() const { return unpack(Packed: Storage); }
170 bool compare_exchange_strong(Data &Existing, Data New);
171
172private:
173 std::atomic<uint64_t> Storage;
174};
175
176/// DataStore record data: 4B + size? + refs? + data + 0
177/// - 4-bytes: Header
178/// - {0,4,8}-bytes: DataSize (may be packed in Header)
179/// - {0,4,8}-bytes: NumRefs (may be packed in Header)
180/// - NumRefs*{4,8}-bytes: Refs[] (end-ptr is 8-byte aligned)
181/// - <data>
182/// - 1-byte: 0-term
183struct DataRecordHandle {
184 /// NumRefs storage: 4B, 2B, 1B, or 0B (no refs). Or, 8B, for alignment
185 /// convenience to avoid computing padding later.
186 enum class NumRefsFlags : uint8_t {
187 Uses0B = 0U,
188 Uses1B = 1U,
189 Uses2B = 2U,
190 Uses4B = 3U,
191 Uses8B = 4U,
192 Max = Uses8B,
193 };
194
195 /// DataSize storage: 8B, 4B, 2B, or 1B.
196 enum class DataSizeFlags {
197 Uses1B = 0U,
198 Uses2B = 1U,
199 Uses4B = 2U,
200 Uses8B = 3U,
201 Max = Uses8B,
202 };
203
204 /// Kind of ref stored in Refs[]: InternalRef or InternalRef4B.
205 enum class RefKindFlags {
206 InternalRef = 0U,
207 InternalRef4B = 1U,
208 Max = InternalRef4B,
209 };
210
211 enum Counts : int {
212 NumRefsShift = 0,
213 NumRefsBits = 3,
214 DataSizeShift = NumRefsShift + NumRefsBits,
215 DataSizeBits = 2,
216 RefKindShift = DataSizeShift + DataSizeBits,
217 RefKindBits = 1,
218 };
219 static_assert(((UINT32_MAX << NumRefsBits) & (uint32_t)NumRefsFlags::Max) ==
220 0,
221 "Not enough bits");
222 static_assert(((UINT32_MAX << DataSizeBits) & (uint32_t)DataSizeFlags::Max) ==
223 0,
224 "Not enough bits");
225 static_assert(((UINT32_MAX << RefKindBits) & (uint32_t)RefKindFlags::Max) ==
226 0,
227 "Not enough bits");
228
229 /// Layout of the DataRecordHandle and how to decode it.
230 struct LayoutFlags {
231 NumRefsFlags NumRefs;
232 DataSizeFlags DataSize;
233 RefKindFlags RefKind;
234
235 static uint64_t pack(LayoutFlags LF) {
236 unsigned Packed = ((unsigned)LF.NumRefs << NumRefsShift) |
237 ((unsigned)LF.DataSize << DataSizeShift) |
238 ((unsigned)LF.RefKind << RefKindShift);
239#ifndef NDEBUG
240 LayoutFlags RoundTrip = unpack(Packed);
241 assert(LF.NumRefs == RoundTrip.NumRefs);
242 assert(LF.DataSize == RoundTrip.DataSize);
243 assert(LF.RefKind == RoundTrip.RefKind);
244#endif
245 return Packed;
246 }
247 static LayoutFlags unpack(uint64_t Storage) {
248 assert(Storage <= UINT8_MAX && "Expect storage to fit in a byte");
249 LayoutFlags LF;
250 LF.NumRefs =
251 (NumRefsFlags)((Storage >> NumRefsShift) & ((1U << NumRefsBits) - 1));
252 LF.DataSize = (DataSizeFlags)((Storage >> DataSizeShift) &
253 ((1U << DataSizeBits) - 1));
254 LF.RefKind =
255 (RefKindFlags)((Storage >> RefKindShift) & ((1U << RefKindBits) - 1));
256 return LF;
257 }
258 };
259
260 /// Header layout:
261 /// - 1-byte: LayoutFlags
262 /// - 1-byte: 1B size field
263 /// - {0,2}-bytes: 2B size field
264 struct Header {
265 using PackTy = uint32_t;
266 PackTy Packed;
267
268 static constexpr unsigned LayoutFlagsShift =
269 (sizeof(PackTy) - 1) * CHAR_BIT;
270 };
271
272 struct Input {
273 InternalRefArrayRef Refs;
274 ArrayRef<char> Data;
275 };
276
277 LayoutFlags getLayoutFlags() const {
278 return LayoutFlags::unpack(Storage: H->Packed >> Header::LayoutFlagsShift);
279 }
280
281 uint64_t getDataSize() const;
282 void skipDataSize(LayoutFlags LF, int64_t &RelOffset) const;
283 uint32_t getNumRefs() const;
284 void skipNumRefs(LayoutFlags LF, int64_t &RelOffset) const;
285 int64_t getRefsRelOffset() const;
286 int64_t getDataRelOffset() const;
287
288 static uint64_t getTotalSize(uint64_t DataRelOffset, uint64_t DataSize) {
289 return DataRelOffset + DataSize + 1;
290 }
291 uint64_t getTotalSize() const {
292 return getDataRelOffset() + getDataSize() + 1;
293 }
294
295 /// Describe the layout of data stored and how to decode from
296 /// DataRecordHandle.
297 struct Layout {
298 explicit Layout(const Input &I);
299
300 LayoutFlags Flags;
301 uint64_t DataSize = 0;
302 uint32_t NumRefs = 0;
303 int64_t RefsRelOffset = 0;
304 int64_t DataRelOffset = 0;
305 uint64_t getTotalSize() const {
306 return DataRecordHandle::getTotalSize(DataRelOffset, DataSize);
307 }
308 };
309
310 InternalRefArrayRef getRefs() const {
311 assert(H && "Expected valid handle");
312 auto *BeginByte = reinterpret_cast<const char *>(H) + getRefsRelOffset();
313 size_t Size = getNumRefs();
314 if (!Size)
315 return InternalRefArrayRef();
316 if (getLayoutFlags().RefKind == RefKindFlags::InternalRef4B)
317 return ArrayRef(reinterpret_cast<const InternalRef4B *>(BeginByte), Size);
318 return ArrayRef(reinterpret_cast<const InternalRef *>(BeginByte), Size);
319 }
320
321 ArrayRef<char> getData() const {
322 assert(H && "Expected valid handle");
323 return ArrayRef(reinterpret_cast<const char *>(H) + getDataRelOffset(),
324 getDataSize());
325 }
326
327 static DataRecordHandle create(function_ref<char *(size_t Size)> Alloc,
328 const Input &I);
329 static Expected<DataRecordHandle>
330 createWithError(function_ref<Expected<char *>(size_t Size)> Alloc,
331 const Input &I);
332 static DataRecordHandle construct(char *Mem, const Input &I);
333
334 static DataRecordHandle get(const char *Mem) {
335 return DataRecordHandle(
336 *reinterpret_cast<const DataRecordHandle::Header *>(Mem));
337 }
338 static Expected<DataRecordHandle>
339 getFromDataPool(const OnDiskDataAllocator &Pool, FileOffset Offset);
340
341 explicit operator bool() const { return H; }
342 const Header &getHeader() const { return *H; }
343
344 DataRecordHandle() = default;
345 explicit DataRecordHandle(const Header &H) : H(&H) {}
346
347private:
348 static DataRecordHandle constructImpl(char *Mem, const Input &I,
349 const Layout &L);
350 const Header *H = nullptr;
351};
352
353/// Proxy for any on-disk object or raw data.
354struct OnDiskContent {
355 std::optional<DataRecordHandle> Record;
356 std::optional<ArrayRef<char>> Bytes;
357
358 ArrayRef<char> getData() const {
359 if (Bytes)
360 return *Bytes;
361 assert(Record && "Expected record or bytes");
362 return Record->getData();
363 }
364};
365
366/// Data loaded inside the memory from standalone file.
367class StandaloneDataInMemory {
368public:
369 OnDiskContent getContent() const;
370
371 OnDiskGraphDB::FileBackedData
372 getInternalFileBackedObjectData(StringRef RootPath) const;
373
374 /// Read this object's data from its file again, so the result does not
375 /// reference \a Region and stays valid after this object is gone.
376 ///
377 /// \returns \c nullptr when it does not apply, and the caller is
378 /// expected to copy instead.
379 std::unique_ptr<MemoryBuffer>
380 getStandaloneMemoryBuffer(StringRef RootPath, StringRef Name,
381 bool RequiresNullTerminator) const;
382
383 StandaloneDataInMemory(std::unique_ptr<sys::fs::mapped_file_region> Region,
384 TrieRecord::StorageKind SK, FileOffset IndexOffset)
385 : Region(std::move(Region)), SK(SK), IndexOffset(IndexOffset) {
386#ifndef NDEBUG
387 bool IsStandalone = false;
388 switch (SK) {
389 case TrieRecord::StorageKind::Standalone:
390 case TrieRecord::StorageKind::StandaloneLeaf:
391 case TrieRecord::StorageKind::StandaloneLeaf0:
392 IsStandalone = true;
393 break;
394 default:
395 break;
396 }
397 assert(IsStandalone);
398#endif
399 }
400
401private:
402 std::unique_ptr<sys::fs::mapped_file_region> Region;
403 TrieRecord::StorageKind SK;
404 FileOffset IndexOffset;
405};
406
407/// Container to lookup loaded standalone objects.
408template <size_t NumShards> class StandaloneDataMap {
409 static_assert(isPowerOf2_64(Value: NumShards), "Expected power of 2");
410
411public:
412 uintptr_t insert(ArrayRef<uint8_t> Hash, TrieRecord::StorageKind SK,
413 std::unique_ptr<sys::fs::mapped_file_region> Region,
414 FileOffset IndexOffset);
415
416 const StandaloneDataInMemory *lookup(ArrayRef<uint8_t> Hash) const;
417 bool count(ArrayRef<uint8_t> Hash) const { return bool(lookup(Hash)); }
418
419private:
420 struct Shard {
421 /// Needs to store a std::unique_ptr for a stable address identity.
422 DenseMap<const uint8_t *, std::unique_ptr<StandaloneDataInMemory>> Map;
423 mutable std::mutex Mutex;
424 };
425 Shard &getShard(ArrayRef<uint8_t> Hash) {
426 return const_cast<Shard &>(
427 const_cast<const StandaloneDataMap *>(this)->getShard(Hash));
428 }
429 const Shard &getShard(ArrayRef<uint8_t> Hash) const {
430 static_assert(NumShards <= 256, "Expected only 8 bits of shard");
431 return Shards[Hash[0] % NumShards];
432 }
433
434 Shard Shards[NumShards];
435};
436
437using StandaloneDataMapTy = StandaloneDataMap<16>;
438
439/// A vector of internal node references.
440class InternalRefVector {
441public:
442 void push_back(InternalRef Ref) {
443 if (NeedsFull)
444 return FullRefs.push_back(Elt: Ref);
445 if (std::optional<InternalRef4B> Small = InternalRef4B::tryToShrink(Ref))
446 return SmallRefs.push_back(Elt: *Small);
447 NeedsFull = true;
448 assert(FullRefs.empty());
449 FullRefs.reserve(N: SmallRefs.size() + 1);
450 for (InternalRef4B Small : SmallRefs)
451 FullRefs.push_back(Elt: Small);
452 FullRefs.push_back(Elt: Ref);
453 SmallRefs.clear();
454 }
455
456 operator InternalRefArrayRef() const {
457 assert(SmallRefs.empty() || FullRefs.empty());
458 return NeedsFull ? InternalRefArrayRef(FullRefs)
459 : InternalRefArrayRef(SmallRefs);
460 }
461
462private:
463 bool NeedsFull = false;
464 SmallVector<InternalRef4B> SmallRefs;
465 SmallVector<InternalRef> FullRefs;
466};
467
468} // namespace
469
470Expected<DataRecordHandle> DataRecordHandle::createWithError(
471 function_ref<Expected<char *>(size_t Size)> Alloc, const Input &I) {
472 Layout L(I);
473 if (Expected<char *> Mem = Alloc(L.getTotalSize()))
474 return constructImpl(Mem: *Mem, I, L);
475 else
476 return Mem.takeError();
477}
478
479ObjectHandle ObjectHandle::fromFileOffset(FileOffset Offset) {
480 // Store the file offset as it is.
481 assert(!(Offset.get() & 0x1));
482 return ObjectHandle(Offset.get());
483}
484
485ObjectHandle ObjectHandle::fromMemory(uintptr_t Ptr) {
486 // Store the pointer from memory with lowest bit set.
487 assert(!(Ptr & 0x1));
488 return ObjectHandle(Ptr | 1);
489}
490
491/// Proxy for an on-disk index record.
492struct OnDiskGraphDB::IndexProxy {
493 FileOffset Offset;
494 ArrayRef<uint8_t> Hash;
495 TrieRecord &Ref;
496};
497
498template <size_t N>
499uintptr_t StandaloneDataMap<N>::insert(
500 ArrayRef<uint8_t> Hash, TrieRecord::StorageKind SK,
501 std::unique_ptr<sys::fs::mapped_file_region> Region,
502 FileOffset IndexOffset) {
503 auto &S = getShard(Hash);
504 std::lock_guard<std::mutex> Lock(S.Mutex);
505 auto &V = S.Map[Hash.data()];
506 if (!V)
507 V = std::make_unique<StandaloneDataInMemory>(args: std::move(Region), args&: SK,
508 args&: IndexOffset);
509 return reinterpret_cast<uintptr_t>(V.get());
510}
511
512template <size_t N>
513const StandaloneDataInMemory *
514StandaloneDataMap<N>::lookup(ArrayRef<uint8_t> Hash) const {
515 auto &S = getShard(Hash);
516 std::lock_guard<std::mutex> Lock(S.Mutex);
517 auto I = S.Map.find(Hash.data());
518 if (I == S.Map.end())
519 return nullptr;
520 return &*I->second;
521}
522
523namespace {
524
525/// Copy of \a sys::fs::TempFile that skips RemoveOnSignal, which is too
526/// expensive to register/unregister at this rate.
527///
528/// FIXME: Add a TempFileManager that maintains a thread-safe list of open temp
529/// files and has a signal handler registerd that removes them all.
530class TempFile {
531 bool Done = false;
532 TempFile(StringRef Name, int FD, OnDiskCASLogger *Logger)
533 : TmpName(std::string(Name)), FD(FD), Logger(Logger) {}
534
535public:
536 /// This creates a temporary file with createUniqueFile.
537 static Expected<TempFile> create(const Twine &Model, OnDiskCASLogger *Logger);
538 TempFile(TempFile &&Other) { *this = std::move(Other); }
539 TempFile &operator=(TempFile &&Other) {
540 TmpName = std::move(Other.TmpName);
541 FD = Other.FD;
542 Logger = Other.Logger;
543 Other.Done = true;
544 Other.FD = -1;
545 return *this;
546 }
547
548 // Name of the temporary file.
549 std::string TmpName;
550
551 // The open file descriptor.
552 int FD = -1;
553
554 OnDiskCASLogger *Logger = nullptr;
555
556 // Keep this with the given name.
557 Error keep(const Twine &Name);
558 Error discard();
559
560 // This checks that keep or delete was called.
561 ~TempFile() { consumeError(Err: discard()); }
562};
563
564class MappedTempFile {
565public:
566 char *data() const { return Map.data(); }
567 size_t size() const { return Map.size(); }
568
569 Error discard() {
570 assert(Map && "Map already destroyed");
571 Map.unmap();
572 return Temp.discard();
573 }
574
575 Error keep(const Twine &Name) {
576 assert(Map && "Map already destroyed");
577 Map.unmap();
578 return Temp.keep(Name);
579 }
580
581 MappedTempFile(TempFile Temp, sys::fs::mapped_file_region Map)
582 : Temp(std::move(Temp)), Map(std::move(Map)) {}
583
584private:
585 TempFile Temp;
586 sys::fs::mapped_file_region Map;
587};
588} // namespace
589
590Error TempFile::discard() {
591 Done = true;
592 if (FD != -1) {
593 sys::fs::file_t File = sys::fs::convertFDToNativeFile(FD);
594 if (std::error_code EC = sys::fs::closeFile(F&: File))
595 return errorCodeToError(EC);
596 }
597 FD = -1;
598
599 // Always try to close and remove.
600 std::error_code RemoveEC;
601 if (!TmpName.empty()) {
602 std::error_code EC = sys::fs::remove(path: TmpName);
603 if (Logger)
604 Logger->logTempFileRemove(TmpName, EC);
605 if (EC)
606 return errorCodeToError(EC);
607 }
608 TmpName = "";
609
610 return Error::success();
611}
612
613Error TempFile::keep(const Twine &Name) {
614 assert(!Done);
615 Done = true;
616 // Always try to close and rename.
617 std::error_code RenameEC = sys::fs::rename(from: TmpName, to: Name);
618
619 if (Logger)
620 Logger->logTempFileKeep(TmpName, Name: Name.str(), EC: RenameEC);
621
622 if (!RenameEC)
623 TmpName = "";
624
625 sys::fs::file_t File = sys::fs::convertFDToNativeFile(FD);
626 if (std::error_code EC = sys::fs::closeFile(F&: File))
627 return errorCodeToError(EC);
628 FD = -1;
629
630 return errorCodeToError(EC: RenameEC);
631}
632
633Expected<TempFile> TempFile::create(const Twine &Model,
634 OnDiskCASLogger *Logger) {
635 int FD;
636 SmallString<128> ResultPath;
637 if (std::error_code EC = sys::fs::createUniqueFile(Model, ResultFD&: FD, ResultPath))
638 return errorCodeToError(EC);
639
640 if (Logger)
641 Logger->logTempFileCreate(Name: ResultPath);
642
643 TempFile Ret(ResultPath, FD, Logger);
644 return std::move(Ret);
645}
646
647bool TrieRecord::compare_exchange_strong(Data &Existing, Data New) {
648 uint64_t ExistingPacked = pack(D: Existing);
649 uint64_t NewPacked = pack(D: New);
650 if (Storage.compare_exchange_strong(i1&: ExistingPacked, i2: NewPacked))
651 return true;
652 Existing = unpack(Packed: ExistingPacked);
653 return false;
654}
655
656Expected<DataRecordHandle>
657DataRecordHandle::getFromDataPool(const OnDiskDataAllocator &Pool,
658 FileOffset Offset) {
659 auto HeaderData = Pool.get(Offset, Size: sizeof(DataRecordHandle::Header));
660 if (!HeaderData)
661 return HeaderData.takeError();
662
663 auto Record = DataRecordHandle::get(Mem: HeaderData->data());
664 if (Record.getTotalSize() + Offset.get() > Pool.size())
665 return createStringError(
666 EC: make_error_code(e: std::errc::illegal_byte_sequence),
667 S: "data record span passed the end of the data pool");
668
669 return Record;
670}
671
672DataRecordHandle DataRecordHandle::constructImpl(char *Mem, const Input &I,
673 const Layout &L) {
674 char *Next = Mem + sizeof(Header);
675
676 // Fill in Packed and set other data, then come back to construct the header.
677 Header::PackTy Packed = 0;
678 Packed |= LayoutFlags::pack(LF: L.Flags) << Header::LayoutFlagsShift;
679
680 // Construct DataSize.
681 switch (L.Flags.DataSize) {
682 case DataSizeFlags::Uses1B:
683 assert(I.Data.size() <= UINT8_MAX);
684 Packed |= (Header::PackTy)I.Data.size()
685 << ((sizeof(Packed) - 2) * CHAR_BIT);
686 break;
687 case DataSizeFlags::Uses2B:
688 assert(I.Data.size() <= UINT16_MAX);
689 Packed |= (Header::PackTy)I.Data.size()
690 << ((sizeof(Packed) - 4) * CHAR_BIT);
691 break;
692 case DataSizeFlags::Uses4B:
693 support::endian::write32le(P: Next, V: I.Data.size());
694 Next += 4;
695 break;
696 case DataSizeFlags::Uses8B:
697 support::endian::write64le(P: Next, V: I.Data.size());
698 Next += 8;
699 break;
700 }
701
702 // Construct NumRefs.
703 //
704 // NOTE: May be writing NumRefs even if there are zero refs in order to fix
705 // alignment.
706 switch (L.Flags.NumRefs) {
707 case NumRefsFlags::Uses0B:
708 break;
709 case NumRefsFlags::Uses1B:
710 assert(I.Refs.size() <= UINT8_MAX);
711 Packed |= (Header::PackTy)I.Refs.size()
712 << ((sizeof(Packed) - 2) * CHAR_BIT);
713 break;
714 case NumRefsFlags::Uses2B:
715 assert(I.Refs.size() <= UINT16_MAX);
716 Packed |= (Header::PackTy)I.Refs.size()
717 << ((sizeof(Packed) - 4) * CHAR_BIT);
718 break;
719 case NumRefsFlags::Uses4B:
720 support::endian::write32le(P: Next, V: I.Refs.size());
721 Next += 4;
722 break;
723 case NumRefsFlags::Uses8B:
724 support::endian::write64le(P: Next, V: I.Refs.size());
725 Next += 8;
726 break;
727 }
728
729 // Construct Refs[].
730 if (!I.Refs.empty()) {
731 assert((L.Flags.RefKind == RefKindFlags::InternalRef4B) == I.Refs.is4B());
732 ArrayRef<uint8_t> RefsBuffer = I.Refs.getBuffer();
733 llvm::copy(Range&: RefsBuffer, Out: Next);
734 Next += RefsBuffer.size();
735 }
736
737 // Construct Data and the trailing null.
738 assert(isAddrAligned(Align(8), Next));
739 llvm::copy(Range: I.Data, Out: Next);
740 Next[I.Data.size()] = 0;
741
742 // Construct the header itself and return.
743 Header *H = new (Mem) Header{.Packed: Packed};
744 DataRecordHandle Record(*H);
745 assert(Record.getData() == I.Data);
746 assert(Record.getNumRefs() == I.Refs.size());
747 assert(Record.getRefs() == I.Refs);
748 assert(Record.getLayoutFlags().DataSize == L.Flags.DataSize);
749 assert(Record.getLayoutFlags().NumRefs == L.Flags.NumRefs);
750 assert(Record.getLayoutFlags().RefKind == L.Flags.RefKind);
751 return Record;
752}
753
754DataRecordHandle::Layout::Layout(const Input &I) {
755 // Start initial relative offsets right after the Header.
756 uint64_t RelOffset = sizeof(Header);
757
758 // Initialize the easy stuff.
759 DataSize = I.Data.size();
760 NumRefs = I.Refs.size();
761
762 // Check refs size.
763 Flags.RefKind =
764 I.Refs.is4B() ? RefKindFlags::InternalRef4B : RefKindFlags::InternalRef;
765
766 // Find the smallest slot available for DataSize.
767 bool Has1B = true;
768 bool Has2B = true;
769 if (DataSize <= UINT8_MAX && Has1B) {
770 Flags.DataSize = DataSizeFlags::Uses1B;
771 Has1B = false;
772 } else if (DataSize <= UINT16_MAX && Has2B) {
773 Flags.DataSize = DataSizeFlags::Uses2B;
774 Has2B = false;
775 } else if (DataSize <= UINT32_MAX) {
776 Flags.DataSize = DataSizeFlags::Uses4B;
777 RelOffset += 4;
778 } else {
779 Flags.DataSize = DataSizeFlags::Uses8B;
780 RelOffset += 8;
781 }
782
783 // Find the smallest slot available for NumRefs. Never sets NumRefs8B here.
784 if (!NumRefs) {
785 Flags.NumRefs = NumRefsFlags::Uses0B;
786 } else if (NumRefs <= UINT8_MAX && Has1B) {
787 Flags.NumRefs = NumRefsFlags::Uses1B;
788 Has1B = false;
789 } else if (NumRefs <= UINT16_MAX && Has2B) {
790 Flags.NumRefs = NumRefsFlags::Uses2B;
791 Has2B = false;
792 } else {
793 Flags.NumRefs = NumRefsFlags::Uses4B;
794 RelOffset += 4;
795 }
796
797 // Helper to "upgrade" either DataSize or NumRefs by 4B to avoid complicated
798 // padding rules when reading and writing. This also bumps RelOffset.
799 //
800 // The value for NumRefs is strictly limited to UINT32_MAX, but it can be
801 // stored as 8B. This means we can *always* find a size to grow.
802 //
803 // NOTE: Only call this once.
804 auto GrowSizeFieldsBy4B = [&]() {
805 assert(isAligned(Align(4), RelOffset));
806 RelOffset += 4;
807
808 assert(Flags.NumRefs != NumRefsFlags::Uses8B &&
809 "Expected to be able to grow NumRefs8B");
810
811 // First try to grow DataSize. NumRefs will not (yet) be 8B, and if
812 // DataSize is upgraded to 8B it'll already be aligned.
813 //
814 // Failing that, grow NumRefs.
815 if (Flags.DataSize < DataSizeFlags::Uses4B)
816 Flags.DataSize = DataSizeFlags::Uses4B; // DataSize: Packed => 4B.
817 else if (Flags.DataSize < DataSizeFlags::Uses8B)
818 Flags.DataSize = DataSizeFlags::Uses8B; // DataSize: 4B => 8B.
819 else if (Flags.NumRefs < NumRefsFlags::Uses4B)
820 Flags.NumRefs = NumRefsFlags::Uses4B; // NumRefs: Packed => 4B.
821 else
822 Flags.NumRefs = NumRefsFlags::Uses8B; // NumRefs: 4B => 8B.
823 };
824
825 assert(isAligned(Align(4), RelOffset));
826 if (Flags.RefKind == RefKindFlags::InternalRef) {
827 // List of 8B refs should be 8B-aligned. Grow one of the sizes to get this
828 // without padding.
829 if (!isAligned(Lhs: Align(8), SizeInBytes: RelOffset))
830 GrowSizeFieldsBy4B();
831
832 assert(isAligned(Align(8), RelOffset));
833 RefsRelOffset = RelOffset;
834 RelOffset += 8 * NumRefs;
835 } else {
836 // The array of 4B refs doesn't need 8B alignment, but the data will need
837 // to be 8B-aligned. Detect this now, and, if necessary, shift everything
838 // by 4B by growing one of the sizes.
839 // If we remove the need for 8B-alignment for data there is <1% savings in
840 // disk storage for a clang build using MCCAS but the 8B-alignment may be
841 // useful in the future so keep it for now.
842 uint64_t RefListSize = 4 * NumRefs;
843 if (!isAligned(Lhs: Align(8), SizeInBytes: RelOffset + RefListSize))
844 GrowSizeFieldsBy4B();
845 RefsRelOffset = RelOffset;
846 RelOffset += RefListSize;
847 }
848
849 assert(isAligned(Align(8), RelOffset));
850 DataRelOffset = RelOffset;
851}
852
853uint64_t DataRecordHandle::getDataSize() const {
854 int64_t RelOffset = sizeof(Header);
855 auto *DataSizePtr = reinterpret_cast<const char *>(H) + RelOffset;
856 switch (getLayoutFlags().DataSize) {
857 case DataSizeFlags::Uses1B:
858 return (H->Packed >> ((sizeof(Header::PackTy) - 2) * CHAR_BIT)) & UINT8_MAX;
859 case DataSizeFlags::Uses2B:
860 return (H->Packed >> ((sizeof(Header::PackTy) - 4) * CHAR_BIT)) &
861 UINT16_MAX;
862 case DataSizeFlags::Uses4B:
863 return support::endian::read32le(P: DataSizePtr);
864 case DataSizeFlags::Uses8B:
865 return support::endian::read64le(P: DataSizePtr);
866 }
867 llvm_unreachable("Unknown DataSizeFlags enum");
868}
869
870void DataRecordHandle::skipDataSize(LayoutFlags LF, int64_t &RelOffset) const {
871 if (LF.DataSize >= DataSizeFlags::Uses4B)
872 RelOffset += 4;
873 if (LF.DataSize >= DataSizeFlags::Uses8B)
874 RelOffset += 4;
875}
876
877uint32_t DataRecordHandle::getNumRefs() const {
878 LayoutFlags LF = getLayoutFlags();
879 int64_t RelOffset = sizeof(Header);
880 skipDataSize(LF, RelOffset);
881 auto *NumRefsPtr = reinterpret_cast<const char *>(H) + RelOffset;
882 switch (LF.NumRefs) {
883 case NumRefsFlags::Uses0B:
884 return 0;
885 case NumRefsFlags::Uses1B:
886 return (H->Packed >> ((sizeof(Header::PackTy) - 2) * CHAR_BIT)) & UINT8_MAX;
887 case NumRefsFlags::Uses2B:
888 return (H->Packed >> ((sizeof(Header::PackTy) - 4) * CHAR_BIT)) &
889 UINT16_MAX;
890 case NumRefsFlags::Uses4B:
891 return support::endian::read32le(P: NumRefsPtr);
892 case NumRefsFlags::Uses8B:
893 return support::endian::read64le(P: NumRefsPtr);
894 }
895 llvm_unreachable("Unknown NumRefsFlags enum");
896}
897
898void DataRecordHandle::skipNumRefs(LayoutFlags LF, int64_t &RelOffset) const {
899 if (LF.NumRefs >= NumRefsFlags::Uses4B)
900 RelOffset += 4;
901 if (LF.NumRefs >= NumRefsFlags::Uses8B)
902 RelOffset += 4;
903}
904
905int64_t DataRecordHandle::getRefsRelOffset() const {
906 LayoutFlags LF = getLayoutFlags();
907 int64_t RelOffset = sizeof(Header);
908 skipDataSize(LF, RelOffset);
909 skipNumRefs(LF, RelOffset);
910 return RelOffset;
911}
912
913int64_t DataRecordHandle::getDataRelOffset() const {
914 LayoutFlags LF = getLayoutFlags();
915 int64_t RelOffset = sizeof(Header);
916 skipDataSize(LF, RelOffset);
917 skipNumRefs(LF, RelOffset);
918 uint32_t RefSize = LF.RefKind == RefKindFlags::InternalRef4B ? 4 : 8;
919 RelOffset += RefSize * getNumRefs();
920 return RelOffset;
921}
922
923Error OnDiskGraphDB::validate(bool Deep, HashingFuncT Hasher) const {
924 if (UpstreamDB) {
925 if (auto E = UpstreamDB->validate(Deep, Hasher))
926 return E;
927 }
928 if (!isAligned(Lhs: Align(8), SizeInBytes: DataPool.size()))
929 return createStringError(EC: llvm::errc::illegal_byte_sequence,
930 S: "data pool bump pointer is not aligned");
931 return Index.validate(RecordVerifier: [&](FileOffset Offset,
932 OnDiskTrieRawHashMap::ConstValueProxy Record)
933 -> Error {
934 auto formatError = [&](Twine Msg) {
935 return createStringError(
936 EC: llvm::errc::illegal_byte_sequence,
937 S: "bad record at 0x" +
938 utohexstr(X: (unsigned)Offset.get(), /*LowerCase=*/true) + ": " +
939 Msg);
940 };
941
942 if (Record.Data.size() != sizeof(TrieRecord))
943 return formatError("wrong data record size");
944 if (!isAligned(Lhs: Align::Of<TrieRecord>(), SizeInBytes: Record.Data.size()))
945 return formatError("wrong data record alignment");
946
947 auto *R = reinterpret_cast<const TrieRecord *>(Record.Data.data());
948 TrieRecord::Data D = R->load();
949 std::unique_ptr<MemoryBuffer> FileBuffer;
950 if ((uint8_t)D.SK != (uint8_t)TrieRecord::StorageKind::Unknown &&
951 (uint8_t)D.SK != (uint8_t)TrieRecord::StorageKind::DataPool &&
952 (uint8_t)D.SK != (uint8_t)TrieRecord::StorageKind::Standalone &&
953 (uint8_t)D.SK != (uint8_t)TrieRecord::StorageKind::StandaloneLeaf &&
954 (uint8_t)D.SK != (uint8_t)TrieRecord::StorageKind::StandaloneLeaf0)
955 return formatError("invalid record kind value");
956
957 auto Ref = InternalRef::getFromOffset(Offset);
958 auto I = getIndexProxyFromRef(Ref);
959 if (!I)
960 return I.takeError();
961
962 switch (D.SK) {
963 case TrieRecord::StorageKind::Unknown:
964 // This could be an abandoned entry due to a termination before updating
965 // the record. It can be reused by later insertion so just skip this entry
966 // for now.
967 return Error::success();
968 case TrieRecord::StorageKind::DataPool: {
969 // Check offset is a postive value, and large enough to hold the
970 // header for the data record.
971 if (D.Offset.get() <= 0 ||
972 D.Offset.get() + sizeof(DataRecordHandle::Header) >= DataPool.size())
973 return formatError("datapool record out of bound");
974
975 // DataRecord start needs to be aligned.
976 if (!isAligned(Lhs: Align(8), SizeInBytes: D.Offset.get()))
977 return formatError("data record offset is not aligned");
978
979 // Validate the layout flags before getFromDataPool calls getTotalSize().
980 auto HeaderData =
981 DataPool.get(Offset: D.Offset, Size: sizeof(DataRecordHandle::Header));
982 if (!HeaderData)
983 return formatError(toString(E: HeaderData.takeError()));
984 auto LF = DataRecordHandle::get(Mem: HeaderData->data()).getLayoutFlags();
985 if (LF.NumRefs > DataRecordHandle::NumRefsFlags::Max ||
986 LF.DataSize > DataRecordHandle::DataSizeFlags::Max)
987 return formatError("data record has invalid layout flags");
988 break;
989 }
990 case TrieRecord::StorageKind::Standalone:
991 case TrieRecord::StorageKind::StandaloneLeaf:
992 case TrieRecord::StorageKind::StandaloneLeaf0:
993 SmallString<256> Path;
994 getStandalonePath(FileSuffix: TrieRecord::getStandaloneFilePrefix(SK: D.SK), IndexOffset: I->Offset,
995 Path);
996 // If need to validate the content of the file later, just load the
997 // buffer here. Otherwise, just check the existance of the file.
998 if (Deep) {
999 auto File = MemoryBuffer::getFile(Filename: Path, /*IsText=*/false,
1000 /*RequiresNullTerminator=*/false);
1001 if (!File || !*File)
1002 return formatError("record file \'" + Path + "\' does not exist");
1003
1004 FileBuffer = std::move(*File);
1005 } else if (!llvm::sys::fs::exists(Path))
1006 return formatError("record file \'" + Path + "\' does not exist");
1007 }
1008
1009 if (!Deep)
1010 return Error::success();
1011
1012 auto dataError = [&](Twine Msg) {
1013 return createStringError(EC: llvm::errc::illegal_byte_sequence,
1014 S: "bad data for digest \'" + toHex(Input: I->Hash) +
1015 "\': " + Msg);
1016 };
1017 SmallVector<ArrayRef<uint8_t>> Refs;
1018 ArrayRef<char> StoredData;
1019
1020 switch (D.SK) {
1021 case TrieRecord::StorageKind::Unknown:
1022 llvm_unreachable("already handled");
1023 case TrieRecord::StorageKind::DataPool: {
1024 auto DataRecord = DataRecordHandle::getFromDataPool(Pool: DataPool, Offset: D.Offset);
1025 if (!DataRecord)
1026 return dataError(toString(E: DataRecord.takeError()));
1027
1028 for (auto InternRef : DataRecord->getRefs()) {
1029 if (InternRef.getFileOffset().get() <= 0)
1030 return dataError("invalid ref offset");
1031 auto Index = getIndexProxyFromRef(Ref: InternRef);
1032 if (!Index)
1033 return Index.takeError();
1034 Refs.push_back(Elt: Index->Hash);
1035 }
1036 StoredData = DataRecord->getData();
1037 break;
1038 }
1039 case TrieRecord::StorageKind::Standalone: {
1040 if (FileBuffer->getBufferSize() < sizeof(DataRecordHandle::Header))
1041 return dataError("data record is not big enough to read the header");
1042 auto DataRecord = DataRecordHandle::get(Mem: FileBuffer->getBufferStart());
1043 if (DataRecord.getTotalSize() < FileBuffer->getBufferSize())
1044 return dataError(
1045 "data record span passed the end of the standalone file");
1046 for (auto InternRef : DataRecord.getRefs()) {
1047 if (InternRef.getFileOffset().get() <= 0)
1048 return dataError("invalid ref offset");
1049 auto Index = getIndexProxyFromRef(Ref: InternRef);
1050 if (!Index)
1051 return Index.takeError();
1052 Refs.push_back(Elt: Index->Hash);
1053 }
1054 StoredData = DataRecord.getData();
1055 break;
1056 }
1057 case TrieRecord::StorageKind::StandaloneLeaf:
1058 case TrieRecord::StorageKind::StandaloneLeaf0: {
1059 StoredData = arrayRefFromStringRef<char>(Input: FileBuffer->getBuffer());
1060 if (D.SK == TrieRecord::StorageKind::StandaloneLeaf0) {
1061 if (!FileBuffer->getBuffer().ends_with(Suffix: '\0'))
1062 return dataError("standalone file is not zero terminated");
1063 StoredData = StoredData.drop_back(N: 1);
1064 }
1065 break;
1066 }
1067 }
1068
1069 SmallVector<uint8_t> ComputedHash;
1070 Hasher(Refs, StoredData, ComputedHash);
1071 if (I->Hash != ArrayRef(ComputedHash))
1072 return dataError("hash mismatch, got \'" + toHex(Input: ComputedHash) +
1073 "\' instead");
1074
1075 return Error::success();
1076 });
1077}
1078
1079Error OnDiskGraphDB::validateObjectID(ObjectID ExternalRef) const {
1080 auto formatError = [&](Twine Msg) {
1081 return createStringError(
1082 EC: llvm::errc::illegal_byte_sequence,
1083 S: "bad ref=0x" +
1084 utohexstr(X: ExternalRef.getOpaqueData(), /*LowerCase=*/true) + ": " +
1085 Msg);
1086 };
1087
1088 if (ExternalRef.getOpaqueData() == 0)
1089 return formatError("zero is not a valid ref");
1090
1091 InternalRef InternalRef = getInternalRef(Ref: ExternalRef);
1092 auto I = getIndexProxyFromRef(Ref: InternalRef);
1093 if (!I)
1094 return formatError(llvm::toString(E: I.takeError()));
1095 auto Hash = getDigest(I: *I);
1096
1097 OnDiskTrieRawHashMap::ConstOnDiskPtr P = Index.find(Hash);
1098 if (!P)
1099 return formatError("not found using hash " + toHex(Input: Hash));
1100 IndexProxy OtherI = getIndexProxyFromPointer(P);
1101 ObjectID OtherRef = getExternalReference(Ref: makeInternalRef(IndexOffset: OtherI.Offset));
1102 if (OtherRef != ExternalRef)
1103 return formatError("ref does not match indexed offset " +
1104 utohexstr(X: OtherRef.getOpaqueData(), /*LowerCase=*/true) +
1105 " for hash " + toHex(Input: Hash));
1106 return Error::success();
1107}
1108
1109void OnDiskGraphDB::print(raw_ostream &OS) const {
1110 OS << "on-disk-root-path: " << RootPath << "\n";
1111
1112 struct PoolInfo {
1113 uint64_t Offset;
1114 };
1115 SmallVector<PoolInfo> Pool;
1116
1117 OS << "\n";
1118 OS << "index:\n";
1119 Index.print(OS, PrintRecordData: [&](ArrayRef<char> Data) {
1120 assert(Data.size() == sizeof(TrieRecord));
1121 assert(isAligned(Align::Of<TrieRecord>(), Data.size()));
1122 auto *R = reinterpret_cast<const TrieRecord *>(Data.data());
1123 TrieRecord::Data D = R->load();
1124 OS << " SK=";
1125 switch (D.SK) {
1126 case TrieRecord::StorageKind::Unknown:
1127 OS << "unknown ";
1128 break;
1129 case TrieRecord::StorageKind::DataPool:
1130 OS << "datapool ";
1131 Pool.push_back(Elt: {.Offset: D.Offset.get()});
1132 break;
1133 case TrieRecord::StorageKind::Standalone:
1134 OS << "standalone-data ";
1135 break;
1136 case TrieRecord::StorageKind::StandaloneLeaf:
1137 OS << "standalone-leaf ";
1138 break;
1139 case TrieRecord::StorageKind::StandaloneLeaf0:
1140 OS << "standalone-leaf+0";
1141 break;
1142 }
1143 OS << " Offset=" << (void *)D.Offset.get();
1144 });
1145 if (Pool.empty())
1146 return;
1147
1148 OS << "\n";
1149 OS << "pool:\n";
1150 llvm::sort(
1151 C&: Pool, Comp: [](PoolInfo LHS, PoolInfo RHS) { return LHS.Offset < RHS.Offset; });
1152 for (PoolInfo PI : Pool) {
1153 OS << "- addr=" << (void *)PI.Offset << " ";
1154 auto D = DataRecordHandle::getFromDataPool(Pool: DataPool, Offset: FileOffset(PI.Offset));
1155 if (!D) {
1156 OS << "error: " << toString(E: D.takeError());
1157 return;
1158 }
1159
1160 OS << "record refs=" << D->getNumRefs() << " data=" << D->getDataSize()
1161 << " size=" << D->getTotalSize()
1162 << " end=" << (void *)(PI.Offset + D->getTotalSize()) << "\n";
1163 }
1164}
1165
1166Expected<OnDiskGraphDB::IndexProxy>
1167OnDiskGraphDB::indexHash(ArrayRef<uint8_t> Hash) {
1168 auto P = Index.insertLazy(
1169 Hash, OnConstruct: [](FileOffset TentativeOffset,
1170 OnDiskTrieRawHashMap::ValueProxy TentativeValue) {
1171 assert(TentativeValue.Data.size() == sizeof(TrieRecord));
1172 assert(
1173 isAddrAligned(Align::Of<TrieRecord>(), TentativeValue.Data.data()));
1174 new (TentativeValue.Data.data()) TrieRecord();
1175 });
1176 if (LLVM_UNLIKELY(!P))
1177 return P.takeError();
1178
1179 assert(*P && "Expected insertion");
1180 return getIndexProxyFromPointer(P: *P);
1181}
1182
1183OnDiskGraphDB::IndexProxy OnDiskGraphDB::getIndexProxyFromPointer(
1184 OnDiskTrieRawHashMap::ConstOnDiskPtr P) const {
1185 assert(P);
1186 assert(P.getOffset());
1187 return IndexProxy{.Offset: P.getOffset(), .Hash: P->Hash,
1188 .Ref: *const_cast<TrieRecord *>(
1189 reinterpret_cast<const TrieRecord *>(P->Data.data()))};
1190}
1191
1192Expected<ObjectID> OnDiskGraphDB::getReference(ArrayRef<uint8_t> Hash) {
1193 auto I = indexHash(Hash);
1194 if (LLVM_UNLIKELY(!I))
1195 return I.takeError();
1196 return getExternalReference(I: *I);
1197}
1198
1199ObjectID OnDiskGraphDB::getExternalReference(const IndexProxy &I) {
1200 return getExternalReference(Ref: makeInternalRef(IndexOffset: I.Offset));
1201}
1202
1203std::optional<ObjectID>
1204OnDiskGraphDB::getExistingReference(ArrayRef<uint8_t> Digest,
1205 bool CheckUpstream) {
1206 auto tryUpstream =
1207 [&](std::optional<IndexProxy> I) -> std::optional<ObjectID> {
1208 if (!CheckUpstream || !UpstreamDB)
1209 return std::nullopt;
1210 std::optional<ObjectID> UpstreamID =
1211 UpstreamDB->getExistingReference(Digest);
1212 if (LLVM_UNLIKELY(!UpstreamID))
1213 return std::nullopt;
1214 auto Ref = expectedToOptional(E: indexHash(Hash: Digest));
1215 if (!Ref)
1216 return std::nullopt;
1217 if (!I)
1218 I.emplace(args&: *Ref);
1219 return getExternalReference(I: *I);
1220 };
1221
1222 OnDiskTrieRawHashMap::ConstOnDiskPtr P = Index.find(Hash: Digest);
1223 if (!P)
1224 return tryUpstream(std::nullopt);
1225 IndexProxy I = getIndexProxyFromPointer(P);
1226 TrieRecord::Data Obj = I.Ref.load();
1227 if (Obj.SK == TrieRecord::StorageKind::Unknown)
1228 return tryUpstream(I);
1229 return getExternalReference(Ref: makeInternalRef(IndexOffset: I.Offset));
1230}
1231
1232Expected<OnDiskGraphDB::IndexProxy>
1233OnDiskGraphDB::getIndexProxyFromRef(InternalRef Ref) const {
1234 auto P = Index.recoverFromFileOffset(Offset: Ref.getFileOffset());
1235 if (LLVM_UNLIKELY(!P))
1236 return P.takeError();
1237 return getIndexProxyFromPointer(P: *P);
1238}
1239
1240Expected<ArrayRef<uint8_t>> OnDiskGraphDB::getDigest(InternalRef Ref) const {
1241 auto I = getIndexProxyFromRef(Ref);
1242 if (!I)
1243 return I.takeError();
1244 return I->Hash;
1245}
1246
1247ArrayRef<uint8_t> OnDiskGraphDB::getDigest(const IndexProxy &I) const {
1248 return I.Hash;
1249}
1250
1251static std::variant<const StandaloneDataInMemory *, DataRecordHandle>
1252getStandaloneDataOrDataRecord(const OnDiskDataAllocator &DataPool,
1253 ObjectHandle OH) {
1254 // Decode ObjectHandle to locate the stored content.
1255 uint64_t Data = OH.getOpaqueData();
1256 if (Data & 1) {
1257 const auto *SDIM =
1258 reinterpret_cast<const StandaloneDataInMemory *>(Data & (-1ULL << 1));
1259 return SDIM;
1260 }
1261
1262 auto DataHandle =
1263 cantFail(ValOrErr: DataRecordHandle::getFromDataPool(Pool: DataPool, Offset: FileOffset(Data)));
1264 assert(DataHandle.getData().end()[0] == 0 && "Null termination");
1265 return DataHandle;
1266}
1267
1268static OnDiskContent getContentFromHandle(const OnDiskDataAllocator &DataPool,
1269 ObjectHandle OH) {
1270 auto SDIMOrRecord = getStandaloneDataOrDataRecord(DataPool, OH);
1271 if (std::holds_alternative<const StandaloneDataInMemory *>(v: SDIMOrRecord)) {
1272 return std::get<const StandaloneDataInMemory *>(v&: SDIMOrRecord)->getContent();
1273 } else {
1274 auto DataHandle = std::get<DataRecordHandle>(v: std::move(SDIMOrRecord));
1275 return OnDiskContent{.Record: std::move(DataHandle), .Bytes: std::nullopt};
1276 }
1277}
1278
1279ArrayRef<char> OnDiskGraphDB::getObjectData(ObjectHandle Node) const {
1280 OnDiskContent Content = getContentFromHandle(DataPool, OH: Node);
1281 return Content.getData();
1282}
1283
1284InternalRefArrayRef OnDiskGraphDB::getInternalRefs(ObjectHandle Node) const {
1285 if (std::optional<DataRecordHandle> Record =
1286 getContentFromHandle(DataPool, OH: Node).Record)
1287 return Record->getRefs();
1288 return std::nullopt;
1289}
1290
1291OnDiskGraphDB::FileBackedData
1292OnDiskGraphDB::getInternalFileBackedObjectData(ObjectHandle Node) const {
1293 auto SDIMOrRecord = getStandaloneDataOrDataRecord(DataPool, OH: Node);
1294 if (std::holds_alternative<const StandaloneDataInMemory *>(v: SDIMOrRecord)) {
1295 auto *SDIM = std::get<const StandaloneDataInMemory *>(v&: SDIMOrRecord);
1296 return SDIM->getInternalFileBackedObjectData(RootPath);
1297 } else {
1298 auto DataHandle = std::get<DataRecordHandle>(v: std::move(SDIMOrRecord));
1299 return FileBackedData{.Data: DataHandle.getData(), /*FileInfo=*/std::nullopt};
1300 }
1301}
1302
1303std::unique_ptr<MemoryBuffer>
1304OnDiskGraphDB::getStandaloneMemoryBuffer(ObjectHandle Node, StringRef Name,
1305 bool RequiresNullTerminator) const {
1306 // Only an object with a file to itself can be read back on its own; one in
1307 // the shared data pool is a subrange of a file holding unrelated objects.
1308 auto SDIMOrRecord = getStandaloneDataOrDataRecord(DataPool, OH: Node);
1309 if (auto **SDIM =
1310 std::get_if<const StandaloneDataInMemory *>(ptr: &SDIMOrRecord)) {
1311 if (std::unique_ptr<MemoryBuffer> Standalone =
1312 (*SDIM)->getStandaloneMemoryBuffer(RootPath, Name,
1313 RequiresNullTerminator))
1314 return Standalone;
1315 }
1316
1317 return MemoryBuffer::getMemBufferCopy(InputData: toStringRef(Input: getObjectData(Node)), BufferName: Name);
1318}
1319
1320Expected<std::optional<ObjectHandle>>
1321OnDiskGraphDB::load(ObjectID ExternalRef) {
1322 InternalRef Ref = getInternalRef(Ref: ExternalRef);
1323 auto I = getIndexProxyFromRef(Ref);
1324 if (!I)
1325 return I.takeError();
1326 TrieRecord::Data Object = I->Ref.load();
1327
1328 if (Object.SK == TrieRecord::StorageKind::Unknown)
1329 return faultInFromUpstream(PrimaryID: ExternalRef);
1330
1331 if (Object.SK == TrieRecord::StorageKind::DataPool)
1332 return ObjectHandle::fromFileOffset(Offset: Object.Offset);
1333
1334 // Only TrieRecord::StorageKind::Standalone (and variants) need to be
1335 // explicitly loaded.
1336 //
1337 // There's corruption if standalone objects have offsets, or if we get here
1338 // for something that isn't standalone.
1339 if (Object.Offset)
1340 return createCorruptObjectError(ID: getDigest(I: *I));
1341 switch (Object.SK) {
1342 case TrieRecord::StorageKind::Unknown:
1343 case TrieRecord::StorageKind::DataPool:
1344 llvm_unreachable("unexpected storage kind");
1345 case TrieRecord::StorageKind::Standalone:
1346 case TrieRecord::StorageKind::StandaloneLeaf0:
1347 case TrieRecord::StorageKind::StandaloneLeaf:
1348 break;
1349 }
1350
1351 // Search in StandaloneMap to see if data is already loaded.
1352 auto *StandaloneMap = static_cast<StandaloneDataMapTy *>(StandaloneData);
1353 if (const StandaloneDataInMemory *SDIM = StandaloneMap->lookup(Hash: I->Hash))
1354 return ObjectHandle::fromMemory(Ptr: reinterpret_cast<uintptr_t>(SDIM));
1355
1356 // Load it from disk.
1357 //
1358 // Note: Creation logic guarantees that data that needs null-termination is
1359 // suitably 0-padded. Requiring null-termination here would be too expensive
1360 // for extremely large objects that happen to be page-aligned.
1361 SmallString<256> Path;
1362 getStandalonePath(FileSuffix: TrieRecord::getStandaloneFilePrefix(SK: Object.SK), IndexOffset: I->Offset,
1363 Path);
1364
1365 auto BypassSandbox = sys::sandbox::scopedDisable();
1366
1367 auto File = sys::fs::openNativeFileForRead(Name: Path);
1368 if (!File)
1369 return createFileError(F: Path, E: File.takeError());
1370
1371 llvm::scope_exit CloseFile([&]() { sys::fs::closeFile(F&: *File); });
1372
1373 sys::fs::file_status Status;
1374 if (std::error_code EC = sys::fs::status(F: *File, Result&: Status))
1375 return createCorruptObjectError(ID: getDigest(I: *I));
1376
1377 std::error_code EC;
1378 auto Region = std::make_unique<sys::fs::mapped_file_region>(
1379 args&: *File, args: sys::fs::mapped_file_region::readonly, args: Status.getSize(), args: 0, args&: EC);
1380 if (EC)
1381 return createCorruptObjectError(ID: getDigest(I: *I));
1382
1383 return ObjectHandle::fromMemory(
1384 Ptr: StandaloneMap->insert(Hash: I->Hash, SK: Object.SK, Region: std::move(Region), IndexOffset: I->Offset));
1385}
1386
1387Expected<bool> OnDiskGraphDB::isMaterialized(ObjectID Ref) {
1388 auto Presence = getObjectPresence(Ref, /*CheckUpstream=*/true);
1389 if (!Presence)
1390 return Presence.takeError();
1391
1392 switch (*Presence) {
1393 case ObjectPresence::Missing:
1394 return false;
1395 case ObjectPresence::InPrimaryDB:
1396 return true;
1397 case ObjectPresence::OnlyInUpstreamDB:
1398 if (auto FaultInResult = faultInFromUpstream(PrimaryID: Ref); !FaultInResult)
1399 return FaultInResult.takeError();
1400 return true;
1401 }
1402 llvm_unreachable("Unknown ObjectPresence enum");
1403}
1404
1405Expected<OnDiskGraphDB::ObjectPresence>
1406OnDiskGraphDB::getObjectPresence(ObjectID ExternalRef,
1407 bool CheckUpstream) const {
1408 InternalRef Ref = getInternalRef(Ref: ExternalRef);
1409 auto I = getIndexProxyFromRef(Ref);
1410 if (!I)
1411 return I.takeError();
1412
1413 TrieRecord::Data Object = I->Ref.load();
1414 if (Object.SK != TrieRecord::StorageKind::Unknown)
1415 return ObjectPresence::InPrimaryDB;
1416
1417 if (!CheckUpstream || !UpstreamDB)
1418 return ObjectPresence::Missing;
1419
1420 std::optional<ObjectID> UpstreamID =
1421 UpstreamDB->getExistingReference(Digest: getDigest(I: *I));
1422 return UpstreamID.has_value() ? ObjectPresence::OnlyInUpstreamDB
1423 : ObjectPresence::Missing;
1424}
1425
1426InternalRef OnDiskGraphDB::makeInternalRef(FileOffset IndexOffset) {
1427 return InternalRef::getFromOffset(Offset: IndexOffset);
1428}
1429
1430static void getStandalonePath(StringRef RootPath, StringRef Prefix,
1431 FileOffset IndexOffset,
1432 SmallVectorImpl<char> &Path) {
1433 Path.assign(in_start: RootPath.begin(), in_end: RootPath.end());
1434 sys::path::append(path&: Path,
1435 a: Prefix + Twine(IndexOffset.get()) + "." + CASFormatVersion);
1436}
1437
1438void OnDiskGraphDB::getStandalonePath(StringRef Prefix, FileOffset IndexOffset,
1439 SmallVectorImpl<char> &Path) const {
1440 return ::getStandalonePath(RootPath, Prefix, IndexOffset, Path);
1441}
1442
1443OnDiskContent StandaloneDataInMemory::getContent() const {
1444 bool Leaf0 = false;
1445 bool Leaf = false;
1446 switch (SK) {
1447 default:
1448 llvm_unreachable("Storage kind must be standalone");
1449 case TrieRecord::StorageKind::Standalone:
1450 break;
1451 case TrieRecord::StorageKind::StandaloneLeaf0:
1452 Leaf = Leaf0 = true;
1453 break;
1454 case TrieRecord::StorageKind::StandaloneLeaf:
1455 Leaf = true;
1456 break;
1457 }
1458
1459 if (Leaf) {
1460 StringRef Data(Region->data(), Region->size());
1461 assert(Data.drop_back(Leaf0).end()[0] == 0 &&
1462 "Standalone node data missing null termination");
1463 return OnDiskContent{.Record: std::nullopt,
1464 .Bytes: arrayRefFromStringRef<char>(Input: Data.drop_back(N: Leaf0))};
1465 }
1466
1467 DataRecordHandle Record = DataRecordHandle::get(Mem: Region->data());
1468 assert(Record.getData().end()[0] == 0 &&
1469 "Standalone object record missing null termination for data");
1470 return OnDiskContent{.Record: Record, .Bytes: std::nullopt};
1471}
1472
1473OnDiskGraphDB::FileBackedData
1474StandaloneDataInMemory::getInternalFileBackedObjectData(
1475 StringRef RootPath) const {
1476 switch (SK) {
1477 case TrieRecord::StorageKind::Unknown:
1478 case TrieRecord::StorageKind::DataPool:
1479 llvm_unreachable("unexpected storage kind");
1480 case TrieRecord::StorageKind::Standalone:
1481 return OnDiskGraphDB::FileBackedData{.Data: getContent().getData(),
1482 /*FileInfo=*/std::nullopt};
1483 case TrieRecord::StorageKind::StandaloneLeaf0:
1484 case TrieRecord::StorageKind::StandaloneLeaf:
1485 bool IsFileNulTerminated = SK == TrieRecord::StorageKind::StandaloneLeaf0;
1486 SmallString<256> Path;
1487 ::getStandalonePath(RootPath, Prefix: TrieRecord::getStandaloneFilePrefix(SK),
1488 IndexOffset, Path);
1489 return OnDiskGraphDB::FileBackedData{
1490 .Data: getContent().getData(), .FileInfo: OnDiskGraphDB::FileBackedData::FileInfoTy{
1491 .FilePath: std::string(Path), .IsFileNulTerminated: IsFileNulTerminated}};
1492 }
1493 llvm_unreachable("Unknown StorageKind enum");
1494}
1495
1496namespace {
1497/// A MemoryBuffer exposing a subrange of another buffer's bytes, under its own
1498/// name.
1499class AdoptedMemoryBuffer final : public MemoryBuffer {
1500public:
1501 AdoptedMemoryBuffer(std::unique_ptr<MemoryBuffer> Buffer, StringRef Name,
1502 uint64_t Offset, uint64_t Size)
1503 : Buffer(std::move(Buffer)), Name(Name.str()) {
1504 const char *Start = this->Buffer->getBufferStart() + Offset;
1505 init(BufStart: Start, BufEnd: Start + Size, /*RequiresNullTerminator=*/false);
1506 }
1507
1508 StringRef getBufferIdentifier() const final { return Name; }
1509
1510 BufferKind getBufferKind() const final { return Buffer->getBufferKind(); }
1511
1512private:
1513 std::unique_ptr<MemoryBuffer> Buffer;
1514 std::string Name;
1515};
1516} // end anonymous namespace
1517
1518std::unique_ptr<MemoryBuffer> StandaloneDataInMemory::getStandaloneMemoryBuffer(
1519 StringRef RootPath, StringRef Name, bool RequiresNullTerminator) const {
1520 // A plain leaf's file is exactly the data, with no nul after it to map. The
1521 // other kinds have one: a record's own terminator, or the one appended to a
1522 // "leaf+0".
1523 if (RequiresNullTerminator && SK == TrieRecord::StorageKind::StandaloneLeaf)
1524 return nullptr;
1525
1526 // Read the file again instead of sharing \a Region, whose lifetime is tied
1527 // to this object. These files are written once and never modified, so the
1528 // second read sees the same bytes. Whether that ends up mapping the file or
1529 // copying it is up to MemoryBuffer; either way the result stands alone.
1530 SmallString<256> Path;
1531 ::getStandalonePath(RootPath, Prefix: TrieRecord::getStandaloneFilePrefix(SK),
1532 IndexOffset, Path);
1533 auto BypassSandbox = sys::sandbox::scopedDisable();
1534 ErrorOr<std::unique_ptr<MemoryBuffer>> Mapped =
1535 MemoryBuffer::getFile(Filename: Path, /*IsText=*/false,
1536 /*RequiresNullTerminator=*/false,
1537 /*IsVolatile=*/false);
1538 if (!Mapped)
1539 return nullptr;
1540
1541 // Find the data within the mapping. A leaf's file holds just the data; a
1542 // record's also holds its header and refs.
1543 OnDiskContent Content = getContent();
1544 ArrayRef<char> Data = Content.getData();
1545 uint64_t Offset = Content.Record ? Data.data() - Region->data() : 0;
1546 if (Offset + Data.size() > (*Mapped)->getBufferSize())
1547 return nullptr;
1548
1549 return std::make_unique<AdoptedMemoryBuffer>(args: std::move(*Mapped), args&: Name, args&: Offset,
1550 args: Data.size());
1551}
1552
1553static Expected<MappedTempFile>
1554createTempFile(StringRef FinalPath, uint64_t Size, OnDiskCASLogger *Logger) {
1555 auto BypassSandbox = sys::sandbox::scopedDisable();
1556
1557 assert(Size && "Unexpected request for an empty temp file");
1558 Expected<TempFile> File = TempFile::create(Model: FinalPath + ".%%%%%%", Logger);
1559 if (!File)
1560 return File.takeError();
1561
1562 if (Error E = preallocateFileTail(FD: File->FD, CurrentSize: 0, NewSize: Size).takeError())
1563 return createFileError(F: File->TmpName, E: std::move(E));
1564
1565 if (auto EC = sys::fs::resize_file_before_mapping_readwrite(FD: File->FD, Size))
1566 return createFileError(F: File->TmpName, EC);
1567
1568 std::error_code EC;
1569 sys::fs::mapped_file_region Map(sys::fs::convertFDToNativeFile(FD: File->FD),
1570 sys::fs::mapped_file_region::readwrite, Size,
1571 0, EC);
1572 if (EC)
1573 return createFileError(F: File->TmpName, EC);
1574 return MappedTempFile(std::move(*File), std::move(Map));
1575}
1576
1577static size_t getPageSize() {
1578 static int PageSize = sys::Process::getPageSizeEstimate();
1579 return PageSize;
1580}
1581
1582Error OnDiskGraphDB::createStandaloneLeaf(IndexProxy &I, ArrayRef<char> Data) {
1583 assert(Data.size() > TrieRecord::MaxEmbeddedSize &&
1584 "Expected a bigger file for external content...");
1585
1586 bool Leaf0 = isAligned(Lhs: Align(getPageSize()), SizeInBytes: Data.size());
1587 TrieRecord::StorageKind SK = Leaf0 ? TrieRecord::StorageKind::StandaloneLeaf0
1588 : TrieRecord::StorageKind::StandaloneLeaf;
1589
1590 SmallString<256> Path;
1591 int64_t FileSize = Data.size() + Leaf0;
1592 getStandalonePath(Prefix: TrieRecord::getStandaloneFilePrefix(SK), IndexOffset: I.Offset, Path);
1593
1594 // Write the file. Don't reuse this mapped_file_region, which is read/write.
1595 // Let load() pull up one that's read-only.
1596 Expected<MappedTempFile> File = createTempFile(FinalPath: Path, Size: FileSize, Logger: Logger.get());
1597 if (!File)
1598 return File.takeError();
1599 assert(File->size() == (uint64_t)FileSize);
1600 llvm::copy(Range&: Data, Out: File->data());
1601 if (Leaf0)
1602 File->data()[Data.size()] = 0;
1603 assert(File->data()[Data.size()] == 0);
1604 if (Error E = File->keep(Name: Path))
1605 return E;
1606
1607 // Store the object reference.
1608 TrieRecord::Data Existing;
1609 {
1610 TrieRecord::Data Leaf{.SK: SK, .Offset: FileOffset()};
1611 if (I.Ref.compare_exchange_strong(Existing, New: Leaf)) {
1612 recordStandaloneSizeIncrease(SizeIncrease: FileSize);
1613 return Error::success();
1614 }
1615 }
1616
1617 // If there was a race, confirm that the new value has valid storage.
1618 if (Existing.SK == TrieRecord::StorageKind::Unknown)
1619 return createCorruptObjectError(ID: getDigest(I));
1620
1621 return Error::success();
1622}
1623
1624Error OnDiskGraphDB::store(ObjectID ID, ArrayRef<ObjectID> Refs,
1625 ArrayRef<char> Data) {
1626 auto I = getIndexProxyFromRef(Ref: getInternalRef(Ref: ID));
1627 if (LLVM_UNLIKELY(!I))
1628 return I.takeError();
1629
1630 // Early return in case the node exists.
1631 {
1632 TrieRecord::Data Existing = I->Ref.load();
1633 if (Existing.SK != TrieRecord::StorageKind::Unknown)
1634 return Error::success();
1635 }
1636
1637 auto BypassSandbox = sys::sandbox::scopedDisable();
1638
1639 // Big leaf nodes.
1640 if (Refs.empty() && Data.size() > TrieRecord::MaxEmbeddedSize)
1641 return createStandaloneLeaf(I&: *I, Data);
1642
1643 // TODO: Check whether it's worth checking the index for an already existing
1644 // object (like storeTreeImpl() does) before building up the
1645 // InternalRefVector.
1646 InternalRefVector InternalRefs;
1647 for (ObjectID Ref : Refs)
1648 InternalRefs.push_back(Ref: getInternalRef(Ref));
1649
1650 // Create the object.
1651
1652 DataRecordHandle::Input Input{.Refs: InternalRefs, .Data: Data};
1653
1654 // Compute the storage kind, allocate it, and create the record.
1655 TrieRecord::StorageKind SK = TrieRecord::StorageKind::Unknown;
1656 FileOffset PoolOffset;
1657 SmallString<256> Path;
1658 std::optional<MappedTempFile> File;
1659 std::optional<uint64_t> FileSize;
1660 auto AllocStandaloneFile = [&](size_t Size) -> Expected<char *> {
1661 getStandalonePath(Prefix: TrieRecord::getStandaloneFilePrefix(
1662 SK: TrieRecord::StorageKind::Standalone),
1663 IndexOffset: I->Offset, Path);
1664 if (Error E = createTempFile(FinalPath: Path, Size, Logger: Logger.get()).moveInto(Value&: File))
1665 return std::move(E);
1666 assert(File->size() == Size);
1667 FileSize = Size;
1668 SK = TrieRecord::StorageKind::Standalone;
1669 return File->data();
1670 };
1671 auto Alloc = [&](size_t Size) -> Expected<char *> {
1672 if (Size <= TrieRecord::MaxEmbeddedSize) {
1673 SK = TrieRecord::StorageKind::DataPool;
1674 auto P = DataPool.allocate(Size);
1675 if (LLVM_UNLIKELY(!P)) {
1676 char *NewAlloc = nullptr;
1677 auto NewE = handleErrors(
1678 E: P.takeError(), Hs: [&](std::unique_ptr<StringError> E) -> Error {
1679 if (E->convertToErrorCode() == std::errc::not_enough_memory)
1680 return AllocStandaloneFile(Size).moveInto(Value&: NewAlloc);
1681 return Error(std::move(E));
1682 });
1683 if (!NewE)
1684 return NewAlloc;
1685 return std::move(NewE);
1686 }
1687 PoolOffset = P->getOffset();
1688 LLVM_DEBUG({
1689 dbgs() << "pool-alloc addr=" << (void *)PoolOffset.get()
1690 << " size=" << Size
1691 << " end=" << (void *)(PoolOffset.get() + Size) << "\n";
1692 });
1693 return (*P)->data();
1694 }
1695 return AllocStandaloneFile(Size);
1696 };
1697
1698 DataRecordHandle Record;
1699 if (Error E =
1700 DataRecordHandle::createWithError(Alloc, I: Input).moveInto(Value&: Record))
1701 return E;
1702 assert(Record.getData().end()[0] == 0 && "Expected null-termination");
1703 assert(Record.getData() == Input.Data && "Expected initialization");
1704 assert(SK != TrieRecord::StorageKind::Unknown);
1705 assert(bool(File) != bool(PoolOffset) &&
1706 "Expected either a mapped file or a pooled offset");
1707
1708 // Check for a race before calling MappedTempFile::keep().
1709 //
1710 // Then decide what to do with the file. Better to discard than overwrite if
1711 // another thread/process has already added this.
1712 TrieRecord::Data Existing = I->Ref.load();
1713 {
1714 TrieRecord::Data NewObject{.SK: SK, .Offset: PoolOffset};
1715 if (File) {
1716 if (Existing.SK == TrieRecord::StorageKind::Unknown) {
1717 // Keep the file!
1718 if (Error E = File->keep(Name: Path))
1719 return E;
1720 } else {
1721 File.reset();
1722 }
1723 }
1724
1725 // If we didn't already see a racing/existing write, then try storing the
1726 // new object. If that races, confirm that the new value has valid storage.
1727 //
1728 // TODO: Find a way to reuse the storage from the new-but-abandoned record
1729 // handle.
1730 if (Existing.SK == TrieRecord::StorageKind::Unknown) {
1731 if (I->Ref.compare_exchange_strong(Existing, New: NewObject)) {
1732 if (FileSize)
1733 recordStandaloneSizeIncrease(SizeIncrease: *FileSize);
1734 return Error::success();
1735 }
1736 }
1737 }
1738
1739 if (Existing.SK == TrieRecord::StorageKind::Unknown)
1740 return createCorruptObjectError(ID: getDigest(I: *I));
1741
1742 // Load existing object.
1743 return Error::success();
1744}
1745
1746Error OnDiskGraphDB::storeFile(ObjectID ID, StringRef FilePath) {
1747 return storeFile(ID, FilePath, /*ImportKind=*/std::nullopt);
1748}
1749
1750Error OnDiskGraphDB::storeFile(
1751 ObjectID ID, StringRef FilePath,
1752 std::optional<InternalUpstreamImportKind> ImportKind) {
1753 auto I = getIndexProxyFromRef(Ref: getInternalRef(Ref: ID));
1754 if (LLVM_UNLIKELY(!I))
1755 return I.takeError();
1756
1757 // Early return in case the node exists.
1758 {
1759 TrieRecord::Data Existing = I->Ref.load();
1760 if (Existing.SK != TrieRecord::StorageKind::Unknown)
1761 return Error::success();
1762 }
1763
1764 auto BypassSandbox = sys::sandbox::scopedDisable();
1765
1766 uint64_t FileSize;
1767 if (std::error_code EC = sys::fs::file_size(Path: FilePath, Result&: FileSize))
1768 return createFileError(F: FilePath, EC);
1769
1770 if (FileSize <= TrieRecord::MaxEmbeddedSize) {
1771 auto Buf = MemoryBuffer::getFile(Filename: FilePath);
1772 if (!Buf)
1773 return createFileError(F: FilePath, EC: Buf.getError());
1774 return store(ID, Refs: {}, Data: arrayRefFromStringRef<char>(Input: (*Buf)->getBuffer()));
1775 }
1776
1777 UniqueTempFile UniqueTmp;
1778 auto ExpectedPath = UniqueTmp.createAndCopyFrom(ParentPath: RootPath, CopyFromPath: FilePath);
1779 if (!ExpectedPath)
1780 return ExpectedPath.takeError();
1781 StringRef TmpPath = *ExpectedPath;
1782
1783 TrieRecord::StorageKind SK;
1784 if (ImportKind.has_value()) {
1785 // Importing the file from upstream, the nul is already added if necessary.
1786 switch (*ImportKind) {
1787 case InternalUpstreamImportKind::Leaf:
1788 SK = TrieRecord::StorageKind::StandaloneLeaf;
1789 break;
1790 case InternalUpstreamImportKind::Leaf0:
1791 SK = TrieRecord::StorageKind::StandaloneLeaf0;
1792 break;
1793 }
1794 } else {
1795 bool Leaf0 = isAligned(Lhs: Align(getPageSize()), SizeInBytes: FileSize);
1796 SK = Leaf0 ? TrieRecord::StorageKind::StandaloneLeaf0
1797 : TrieRecord::StorageKind::StandaloneLeaf;
1798
1799 if (Leaf0) {
1800 // Add a nul byte at the end.
1801 std::error_code EC;
1802 raw_fd_ostream OS(TmpPath, EC, sys::fs::CD_OpenExisting,
1803 sys::fs::FA_Write, sys::fs::OF_Append);
1804 if (EC)
1805 return createFileError(F: TmpPath, EC);
1806 OS.write(C: 0);
1807 OS.close();
1808 if (OS.has_error())
1809 return createFileError(F: TmpPath, EC: OS.error());
1810 }
1811 }
1812
1813 SmallString<256> StandalonePath;
1814 getStandalonePath(Prefix: TrieRecord::getStandaloneFilePrefix(SK), IndexOffset: I->Offset,
1815 Path&: StandalonePath);
1816 if (Error E = UniqueTmp.renameTo(RenameToPath: StandalonePath))
1817 return E;
1818
1819 // Store the object reference.
1820 TrieRecord::Data Existing;
1821 {
1822 TrieRecord::Data Leaf{.SK: SK, .Offset: FileOffset()};
1823 if (I->Ref.compare_exchange_strong(Existing, New: Leaf)) {
1824 recordStandaloneSizeIncrease(SizeIncrease: FileSize);
1825 return Error::success();
1826 }
1827 }
1828
1829 // If there was a race, confirm that the new value has valid storage.
1830 if (Existing.SK == TrieRecord::StorageKind::Unknown)
1831 return createCorruptObjectError(ID: getDigest(I: *I));
1832
1833 return Error::success();
1834}
1835
1836void OnDiskGraphDB::recordStandaloneSizeIncrease(size_t SizeIncrease) {
1837 standaloneStorageSize().fetch_add(i: SizeIncrease, m: std::memory_order_relaxed);
1838}
1839
1840std::atomic<uint64_t> &OnDiskGraphDB::standaloneStorageSize() const {
1841 MutableArrayRef<uint8_t> UserHeader = DataPool.getUserHeader();
1842 assert(UserHeader.size() == sizeof(std::atomic<uint64_t>));
1843 assert(isAddrAligned(Align(8), UserHeader.data()));
1844 return *reinterpret_cast<std::atomic<uint64_t> *>(UserHeader.data());
1845}
1846
1847uint64_t OnDiskGraphDB::getStandaloneStorageSize() const {
1848 return standaloneStorageSize().load(m: std::memory_order_relaxed);
1849}
1850
1851size_t OnDiskGraphDB::getStorageSize() const {
1852 return Index.size() + DataPool.size() + getStandaloneStorageSize();
1853}
1854
1855unsigned OnDiskGraphDB::getHardStorageLimitUtilization() const {
1856 unsigned IndexPercent = Index.size() * 100ULL / Index.capacity();
1857 unsigned DataPercent = DataPool.size() * 100ULL / DataPool.capacity();
1858 return std::max(a: IndexPercent, b: DataPercent);
1859}
1860
1861Expected<std::unique_ptr<OnDiskGraphDB>>
1862OnDiskGraphDB::open(StringRef AbsPath, StringRef HashName,
1863 unsigned HashByteSize, OnDiskGraphDB *UpstreamDB,
1864 std::shared_ptr<OnDiskCASLogger> Logger,
1865 FaultInPolicy Policy) {
1866 if (std::error_code EC = sys::fs::create_directories(path: AbsPath))
1867 return createFileError(F: AbsPath, EC);
1868
1869 constexpr uint64_t MB = 1024ull * 1024ull;
1870 constexpr uint64_t GB = 1024ull * 1024ull * 1024ull;
1871
1872 uint64_t MaxIndexSize = 12 * GB;
1873 uint64_t MaxDataPoolSize = 24 * GB;
1874
1875 if (useSmallMappingSize(Path: AbsPath)) {
1876 MaxIndexSize = 1 * GB;
1877 MaxDataPoolSize = 2 * GB;
1878 }
1879
1880 auto CustomSize = getOverriddenMaxMappingSize();
1881 if (!CustomSize)
1882 return CustomSize.takeError();
1883 if (*CustomSize)
1884 MaxIndexSize = MaxDataPoolSize = **CustomSize;
1885
1886 SmallString<256> IndexPath(AbsPath);
1887 sys::path::append(path&: IndexPath, a: IndexFilePrefix + CASFormatVersion);
1888 std::optional<OnDiskTrieRawHashMap> Index;
1889 if (Error E = OnDiskTrieRawHashMap::create(
1890 Path: IndexPath, TrieName: IndexTableName + "[" + HashName + "]",
1891 NumHashBits: HashByteSize * CHAR_BIT,
1892 /*DataSize=*/sizeof(TrieRecord), MaxFileSize: MaxIndexSize,
1893 /*MinFileSize=*/NewFileInitialSize: MB, Logger)
1894 .moveInto(Value&: Index))
1895 return std::move(E);
1896
1897 uint32_t UserHeaderSize = sizeof(std::atomic<uint64_t>);
1898
1899 SmallString<256> DataPoolPath(AbsPath);
1900 sys::path::append(path&: DataPoolPath, a: DataPoolFilePrefix + CASFormatVersion);
1901 std::optional<OnDiskDataAllocator> DataPool;
1902 StringRef PolicyName =
1903 Policy == FaultInPolicy::SingleNode ? "single" : "full";
1904 if (Error E = OnDiskDataAllocator::create(
1905 Path: DataPoolPath,
1906 TableName: DataPoolTableName + "[" + HashName + "]" + PolicyName,
1907 MaxFileSize: MaxDataPoolSize, /*MinFileSize=*/NewFileInitialSize: MB, UserHeaderSize, Logger,
1908 UserHeaderInit: [](void *UserHeaderPtr) {
1909 new (UserHeaderPtr) std::atomic<uint64_t>(0);
1910 })
1911 .moveInto(Value&: DataPool))
1912 return std::move(E);
1913 if (DataPool->getUserHeader().size() != UserHeaderSize)
1914 return createStringError(EC: llvm::errc::argument_out_of_domain,
1915 S: "unexpected user header in '" + DataPoolPath +
1916 "'");
1917
1918 return std::unique_ptr<OnDiskGraphDB>(
1919 new OnDiskGraphDB(AbsPath, std::move(*Index), std::move(*DataPool),
1920 UpstreamDB, Policy, std::move(Logger)));
1921}
1922
1923OnDiskGraphDB::OnDiskGraphDB(StringRef RootPath, OnDiskTrieRawHashMap Index,
1924 OnDiskDataAllocator DataPool,
1925 OnDiskGraphDB *UpstreamDB, FaultInPolicy Policy,
1926 std::shared_ptr<OnDiskCASLogger> Logger)
1927 : Index(std::move(Index)), DataPool(std::move(DataPool)),
1928 RootPath(RootPath.str()), UpstreamDB(UpstreamDB), FIPolicy(Policy),
1929 Logger(std::move(Logger)) {
1930 /// Lifetime for "big" objects not in DataPool.
1931 ///
1932 /// NOTE: Could use ThreadSafeTrieRawHashMap here. For now, doing something
1933 /// simpler on the assumption there won't be much contention since most data
1934 /// is not big. If there is contention, and we've already fixed ObjectProxy
1935 /// object handles to be cheap enough to use consistently, the fix might be
1936 /// to use better use of them rather than optimizing this map.
1937 ///
1938 /// FIXME: Figure out the right number of shards, if any.
1939 StandaloneData = new StandaloneDataMapTy();
1940}
1941
1942OnDiskGraphDB::~OnDiskGraphDB() {
1943 delete static_cast<StandaloneDataMapTy *>(StandaloneData);
1944}
1945
1946Error OnDiskGraphDB::importFullTree(ObjectID PrimaryID,
1947 ObjectHandle UpstreamNode) {
1948 // Copies the full CAS tree from upstream. Uses depth-first copying to protect
1949 // against the process dying during importing and leaving the database with an
1950 // incomplete tree. Note that if the upstream has missing nodes then the tree
1951 // will be copied with missing nodes as well, it won't be considered an error.
1952 struct UpstreamCursor {
1953 ObjectHandle Node;
1954 size_t RefsCount;
1955 object_refs_iterator RefI;
1956 object_refs_iterator RefE;
1957 };
1958 /// Keeps track of the state of visitation for current node and all of its
1959 /// parents.
1960 SmallVector<UpstreamCursor, 16> CursorStack;
1961 /// Keeps track of the currently visited nodes as they are imported into
1962 /// primary database, from current node and its parents. When a node is
1963 /// entered for visitation it appends its own ID, then appends referenced IDs
1964 /// as they get imported. When a node is fully imported it removes the
1965 /// referenced IDs from the bottom of the stack which leaves its own ID at the
1966 /// bottom, adding to the list of referenced IDs for the parent node.
1967 SmallVector<ObjectID, 128> PrimaryNodesStack;
1968
1969 auto enqueueNode = [&](ObjectID PrimaryID, std::optional<ObjectHandle> Node) {
1970 PrimaryNodesStack.push_back(Elt: PrimaryID);
1971 if (!Node)
1972 return;
1973 auto Refs = UpstreamDB->getObjectRefs(Node: *Node);
1974 CursorStack.push_back(
1975 Elt: {.Node: *Node, .RefsCount: (size_t)llvm::size(Range&: Refs), .RefI: Refs.begin(), .RefE: Refs.end()});
1976 };
1977
1978 enqueueNode(PrimaryID, UpstreamNode);
1979
1980 while (!CursorStack.empty()) {
1981 UpstreamCursor &Cur = CursorStack.back();
1982 if (Cur.RefI == Cur.RefE) {
1983 // Copy the node data into the primary store.
1984
1985 // The bottom of \p PrimaryNodesStack contains the primary ID for the
1986 // current node plus the list of imported referenced IDs.
1987 assert(PrimaryNodesStack.size() >= Cur.RefsCount + 1);
1988 ObjectID PrimaryID = *(PrimaryNodesStack.end() - Cur.RefsCount - 1);
1989 auto PrimaryRefs = ArrayRef(PrimaryNodesStack)
1990 .slice(N: PrimaryNodesStack.size() - Cur.RefsCount);
1991 if (Error E = importUpstreamData(PrimaryID, PrimaryRefs, UpstreamNode: Cur.Node))
1992 return E;
1993 // Remove the current node and its IDs from the stack.
1994 PrimaryNodesStack.truncate(N: PrimaryNodesStack.size() - Cur.RefsCount);
1995 CursorStack.pop_back();
1996 continue;
1997 }
1998
1999 ObjectID UpstreamID = *(Cur.RefI++);
2000 auto PrimaryID = getReference(Hash: UpstreamDB->getDigest(Ref: UpstreamID));
2001 if (LLVM_UNLIKELY(!PrimaryID))
2002 return PrimaryID.takeError();
2003 if (containsObject(Ref: *PrimaryID, /*CheckUpstream=*/false)) {
2004 // This \p ObjectID already exists in the primary. Either it was imported
2005 // via \p importFullTree or the client created it, in which case the
2006 // client takes responsibility for how it was formed.
2007 enqueueNode(*PrimaryID, std::nullopt);
2008 continue;
2009 }
2010 Expected<std::optional<ObjectHandle>> UpstreamNode =
2011 UpstreamDB->load(ExternalRef: UpstreamID);
2012 if (!UpstreamNode)
2013 return UpstreamNode.takeError();
2014 enqueueNode(*PrimaryID, *UpstreamNode);
2015 }
2016
2017 assert(PrimaryNodesStack.size() == 1);
2018 assert(PrimaryNodesStack.front() == PrimaryID);
2019 return Error::success();
2020}
2021
2022Error OnDiskGraphDB::importSingleNode(ObjectID PrimaryID,
2023 ObjectHandle UpstreamNode) {
2024 // Copies only a single node, it doesn't copy the referenced nodes.
2025
2026 auto UpstreamRefs = UpstreamDB->getObjectRefs(Node: UpstreamNode);
2027 SmallVector<ObjectID, 64> Refs;
2028 Refs.reserve(N: llvm::size(Range&: UpstreamRefs));
2029 for (ObjectID UpstreamRef : UpstreamRefs) {
2030 auto Ref = getReference(Hash: UpstreamDB->getDigest(Ref: UpstreamRef));
2031 if (LLVM_UNLIKELY(!Ref))
2032 return Ref.takeError();
2033 Refs.push_back(Elt: *Ref);
2034 }
2035
2036 return importUpstreamData(PrimaryID, PrimaryRefs: Refs, UpstreamNode);
2037}
2038
2039Error OnDiskGraphDB::importUpstreamData(ObjectID PrimaryID,
2040 ArrayRef<ObjectID> PrimaryRefs,
2041 ObjectHandle UpstreamNode) {
2042 // If there are references we can't copy an upstream's standalone file because
2043 // we need to re-resolve the reference offsets it contains.
2044 if (PrimaryRefs.empty()) {
2045 auto FBData = UpstreamDB->getInternalFileBackedObjectData(Node: UpstreamNode);
2046 if (FBData.FileInfo.has_value()) {
2047 // Disk-space optimization, import the file directly since it is a
2048 // standalone leaf.
2049 return storeFile(
2050 ID: PrimaryID, FilePath: FBData.FileInfo->FilePath,
2051 /*InternalUpstreamImport=*/ImportKind: FBData.FileInfo->IsFileNulTerminated
2052 ? InternalUpstreamImportKind::Leaf0
2053 : InternalUpstreamImportKind::Leaf);
2054 }
2055 }
2056
2057 auto Data = UpstreamDB->getObjectData(Node: UpstreamNode);
2058 return store(ID: PrimaryID, Refs: PrimaryRefs, Data);
2059}
2060
2061Expected<std::optional<ObjectHandle>>
2062OnDiskGraphDB::faultInFromUpstream(ObjectID PrimaryID) {
2063 if (!UpstreamDB)
2064 return std::nullopt;
2065
2066 auto UpstreamID = UpstreamDB->getReference(Hash: getDigest(Ref: PrimaryID));
2067 if (LLVM_UNLIKELY(!UpstreamID))
2068 return UpstreamID.takeError();
2069
2070 Expected<std::optional<ObjectHandle>> UpstreamNode =
2071 UpstreamDB->load(ExternalRef: *UpstreamID);
2072 if (!UpstreamNode)
2073 return UpstreamNode.takeError();
2074 if (!*UpstreamNode)
2075 return std::nullopt;
2076
2077 if (Error E = FIPolicy == FaultInPolicy::SingleNode
2078 ? importSingleNode(PrimaryID, UpstreamNode: **UpstreamNode)
2079 : importFullTree(PrimaryID, UpstreamNode: **UpstreamNode))
2080 return std::move(E);
2081 return load(ExternalRef: PrimaryID);
2082}
2083