startupcache/StartupCache.cpp

   1 /* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
   2 /* vim: set ts=8 sts=2 et sw=2 tw=80: */
   3 /* This Source Code Form is subject to the terms of the Mozilla Public
   4  * License, v. 2.0. If a copy of the MPL was not distributed with this
   5  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
   6
   7 #include "prio.h"
   8 #include "PLDHashTable.h"
   9 #include "mozilla/IOInterposer.h"
  10 #include "mozilla/AutoMemMap.h"
  11 #include "mozilla/IOBuffers.h"
  12 #include "mozilla/MemoryReporting.h"
  13 #include "mozilla/MemUtils.h"
  14 #include "mozilla/MmapFaultHandler.h"
  15 #include "mozilla/ResultExtensions.h"
  16 #include "mozilla/scache/StartupCache.h"
  17 #include "mozilla/ScopeExit.h"
  18
  19 #include "nsClassHashtable.h"
  20 #include "nsComponentManagerUtils.h"
  21 #include "nsCRT.h"
  22 #include "nsDirectoryServiceUtils.h"
  23 #include "nsIClassInfo.h"
  24 #include "nsIFile.h"
  25 #include "nsIObserver.h"
  26 #include "nsIOutputStream.h"
  27 #include "nsISupports.h"
  28 #include "nsITimer.h"
  29 #include "mozilla/Omnijar.h"
  30 #include "prenv.h"
  31 #include "mozilla/Telemetry.h"
  32 #include "nsThreadUtils.h"
  33 #include "nsXULAppAPI.h"
  34 #include "nsIProtocolHandler.h"
  35 #include "GeckoProfiler.h"
  36 #include "nsAppRunner.h"
  37 #include "xpcpublic.h"
  38 #ifdef MOZ_BACKGROUNDTASKS
  39 #  include "mozilla/BackgroundTasks.h"
  40 #endif
  41
  42 #if defined(XP_WIN)
  43 #  include <windows.h>
  44 #endif
  45
  46 #ifdef IS_BIG_ENDIAN
  47 #  define SC_ENDIAN "big"
  48 #else
  49 #  define SC_ENDIAN "little"
  50 #endif
  51
  52 #if PR_BYTES_PER_WORD == 4
  53 #  define SC_WORDSIZE "4"
  54 #else
  55 #  define SC_WORDSIZE "8"
  56 #endif
  57
  58 using namespace mozilla::Compression;
  59
  60 namespace mozilla {
  61 namespace scache {
  62
  63 MOZ_DEFINE_MALLOC_SIZE_OF(StartupCacheMallocSizeOf)
  64
  65 NS_IMETHODIMP
  66 StartupCache::CollectReports(nsIHandleReportCallback* aHandleReport,
  67                              nsISupports* aData, bool aAnonymize) {
  68   MOZ_COLLECT_REPORT(
  69       "explicit/startup-cache/mapping", KIND_NONHEAP, UNITS_BYTES,
  70       mCacheData.nonHeapSizeOfExcludingThis(),
  71       "Memory used to hold the mapping of the startup cache from file. "
  72       "This memory is likely to be swapped out shortly after start-up.");
  73
  74   MOZ_COLLECT_REPORT("explicit/startup-cache/data", KIND_HEAP, UNITS_BYTES,
  75                      HeapSizeOfIncludingThis(StartupCacheMallocSizeOf),
  76                      "Memory used by the startup cache for things other than "
  77                      "the file mapping.");
  78
  79   return NS_OK;
  80 }
  81
  82 static const uint8_t MAGIC[] = "startupcache0002";
  83 // This is a heuristic value for how much to reserve for mTable to avoid
  84 // rehashing. This is not a hard limit in release builds, but it is in
  85 // debug builds as it should be stable. If we exceed this number we should
  86 // just increase it.
  87 static const size_t STARTUP_CACHE_RESERVE_CAPACITY = 450;
  88 // This is a hard limit which we will assert on, to ensure that we don't
  89 // have some bug causing runaway cache growth.
  90 static const size_t STARTUP_CACHE_MAX_CAPACITY = 5000;
  91
  92 // Not const because we change it for gtests.
  93 static uint8_t STARTUP_CACHE_WRITE_TIMEOUT = 60;
  94
  95 #define STARTUP_CACHE_NAME "startupCache." SC_WORDSIZE "." SC_ENDIAN
  96
  97 static inline Result<Ok, nsresult> Write(PRFileDesc* fd, const void* data,
  98                                          int32_t len) {
  99   if (PR_Write(fd, data, len) != len) {
 100     return Err(NS_ERROR_FAILURE);
 101   }
 102   return Ok();
 103 }
 104
 105 static inline Result<Ok, nsresult> Seek(PRFileDesc* fd, int32_t offset) {
 106   if (PR_Seek(fd, offset, PR_SEEK_SET) == -1) {
 107     return Err(NS_ERROR_FAILURE);
 108   }
 109   return Ok();
 110 }
 111
 112 static nsresult MapLZ4ErrorToNsresult(size_t aError) {
 113   return NS_ERROR_FAILURE;
 114 }
 115
 116 StartupCache* StartupCache::GetSingletonNoInit() {
 117   return StartupCache::gStartupCache;
 118 }
 119
 120 StartupCache* StartupCache::GetSingleton() {
 121 #ifdef MOZ_BACKGROUNDTASKS
 122   if (BackgroundTasks::IsBackgroundTaskMode()) {
 123     return nullptr;
 124   }
 125 #endif
 126
 127   if (!gStartupCache) {
 128     if (!XRE_IsParentProcess()) {
 129       return nullptr;
 130     }
 131 #ifdef MOZ_DISABLE_STARTUPCACHE
 132     return nullptr;
 133 #else
 134     StartupCache::InitSingleton();
 135 #endif
 136   }
 137
 138   return StartupCache::gStartupCache;
 139 }
 140
 141 void StartupCache::DeleteSingleton() { StartupCache::gStartupCache = nullptr; }
 142
 143 nsresult StartupCache::InitSingleton() {
 144   nsresult rv;
 145   StartupCache::gStartupCache = new StartupCache();
 146
 147   rv = StartupCache::gStartupCache->Init();
 148   if (NS_FAILED(rv)) {
 149     StartupCache::gStartupCache = nullptr;
 150   }
 151   return rv;
 152 }
 153
 154 StaticRefPtr<StartupCache> StartupCache::gStartupCache;
 155 bool StartupCache::gShutdownInitiated;
 156 bool StartupCache::gIgnoreDiskCache;
 157 bool StartupCache::gFoundDiskCacheOnInit;
 158
 159 NS_IMPL_ISUPPORTS(StartupCache, nsIMemoryReporter)
 160
 161 StartupCache::StartupCache()
 162     : mTableLock("StartupCache::mTableLock"),
 163       mDirty(false),
 164       mWrittenOnce(false),
 165       mCurTableReferenced(false),
 166       mRequestedCount(0),
 167       mCacheEntriesBaseOffset(0),
 168       mPrefetchThread(nullptr) {}
 169
 170 StartupCache::~StartupCache() { UnregisterWeakMemoryReporter(this); }
 171
 172 nsresult StartupCache::Init() {
 173   // workaround for bug 653936
 174   nsCOMPtr<nsIProtocolHandler> jarInitializer(
 175       do_GetService(NS_NETWORK_PROTOCOL_CONTRACTID_PREFIX "jar"));
 176
 177   nsresult rv;
 178
 179   if (mozilla::RunningGTest()) {
 180     STARTUP_CACHE_WRITE_TIMEOUT = 3;
 181   }
 182
 183   // This allows to override the startup cache filename
 184   // which is useful from xpcshell, when there is no ProfLDS directory to keep
 185   // cache in.
 186   char* env = PR_GetEnv("MOZ_STARTUP_CACHE");
 187   if (env && *env) {
 188     rv = NS_NewLocalFile(NS_ConvertUTF8toUTF16(env), false,
 189                          getter_AddRefs(mFile));
 190   } else {
 191     nsCOMPtr<nsIFile> file;
 192     rv = NS_GetSpecialDirectory("ProfLDS", getter_AddRefs(file));
 193     if (NS_FAILED(rv)) {
 194       // return silently, this will fail in mochitests's xpcshell process.
 195       return rv;
 196     }
 197
 198     rv = file->AppendNative("startupCache"_ns);
 199     NS_ENSURE_SUCCESS(rv, rv);
 200
 201     // Try to create the directory if it's not there yet
 202     rv = file->Create(nsIFile::DIRECTORY_TYPE, 0777);
 203     if (NS_FAILED(rv) && rv != NS_ERROR_FILE_ALREADY_EXISTS) return rv;
 204
 205     rv = file->AppendNative(nsLiteralCString(STARTUP_CACHE_NAME));
 206
 207     NS_ENSURE_SUCCESS(rv, rv);
 208
 209     mFile = file;
 210   }
 211
 212   NS_ENSURE_TRUE(mFile, NS_ERROR_UNEXPECTED);
 213
 214   mObserverService = do_GetService("@mozilla.org/observer-service;1");
 215
 216   if (!mObserverService) {
 217     NS_WARNING("Could not get observerService.");
 218     return NS_ERROR_UNEXPECTED;
 219   }
 220
 221   mListener = new StartupCacheListener();
 222   rv = mObserverService->AddObserver(mListener, NS_XPCOM_SHUTDOWN_OBSERVER_ID,
 223                                      false);
 224   NS_ENSURE_SUCCESS(rv, rv);
 225   rv = mObserverService->AddObserver(mListener, "startupcache-invalidate",
 226                                      false);
 227   NS_ENSURE_SUCCESS(rv, rv);
 228
 229   auto result = LoadArchive();
 230   rv = result.isErr() ? result.unwrapErr() : NS_OK;
 231
 232   gFoundDiskCacheOnInit = rv != NS_ERROR_FILE_NOT_FOUND;
 233
 234   // Sometimes we don't have a cache yet, that's ok.
 235   // If it's corrupted, just remove it and start over.
 236   if (gIgnoreDiskCache || (NS_FAILED(rv) && rv != NS_ERROR_FILE_NOT_FOUND)) {
 237     NS_WARNING("Failed to load startupcache file correctly, removing!");
 238     InvalidateCache();
 239   }
 240
 241   RegisterWeakMemoryReporter(this);
 242   mDecompressionContext = MakeUnique<LZ4FrameDecompressionContext>(true);
 243
 244   return NS_OK;
 245 }
 246
 247 void StartupCache::StartPrefetchMemoryThread() {
 248   // XXX: It would be great for this to not create its own thread, unfortunately
 249   // there doesn't seem to be an existing thread that makes sense for this, so
 250   // barring a coordinated global scheduling system this is the best we get.
 251   mPrefetchThread = PR_CreateThread(
 252       PR_USER_THREAD, StartupCache::ThreadedPrefetch, this, PR_PRIORITY_NORMAL,
 253       PR_GLOBAL_THREAD, PR_JOINABLE_THREAD, 256 * 1024);
 254 }
 255
 256 /**
 257  * LoadArchive can only be called from the main thread.
 258  */
 259 Result<Ok, nsresult> StartupCache::LoadArchive() {
 260   MOZ_ASSERT(NS_IsMainThread(), "Can only load startup cache on main thread");
 261   if (gIgnoreDiskCache) return Err(NS_ERROR_FAILURE);
 262
 263   MOZ_TRY(mCacheData.init(mFile));
 264   auto size = mCacheData.size();
 265   if (CanPrefetchMemory()) {
 266     StartPrefetchMemoryThread();
 267   }
 268
 269   uint32_t headerSize;
 270   if (size < sizeof(MAGIC) + sizeof(headerSize)) {
 271     return Err(NS_ERROR_UNEXPECTED);
 272   }
 273
 274   auto data = mCacheData.get<uint8_t>();
 275   auto end = data + size;
 276
 277   MMAP_FAULT_HANDLER_BEGIN_BUFFER(data.get(), size)
 278
 279   if (memcmp(MAGIC, data.get(), sizeof(MAGIC))) {
 280     return Err(NS_ERROR_UNEXPECTED);
 281   }
 282   data += sizeof(MAGIC);
 283
 284   headerSize = LittleEndian::readUint32(data.get());
 285   data += sizeof(headerSize);
 286
 287   if (headerSize > end - data) {
 288     MOZ_ASSERT(false, "StartupCache file is corrupt.");
 289     return Err(NS_ERROR_UNEXPECTED);
 290   }
 291
 292   Range<uint8_t> header(data, data + headerSize);
 293   data += headerSize;
 294
 295   mCacheEntriesBaseOffset = sizeof(MAGIC) + sizeof(headerSize) + headerSize;
 296   {
 297     if (!mTable.reserve(STARTUP_CACHE_RESERVE_CAPACITY)) {
 298       return Err(NS_ERROR_UNEXPECTED);
 299     }
 300     auto cleanup = MakeScopeExit([&]() {
 301       WaitOnPrefetchThread();
 302       mTable.clear();
 303       mCacheData.reset();
 304     });
 305     loader::InputBuffer buf(header);
 306
 307     uint32_t currentOffset = 0;
 308     while (!buf.finished()) {
 309       uint32_t offset = 0;
 310       uint32_t compressedSize = 0;
 311       uint32_t uncompressedSize = 0;
 312       nsCString key;
 313       buf.codeUint32(offset);
 314       buf.codeUint32(compressedSize);
 315       buf.codeUint32(uncompressedSize);
 316       buf.codeString(key);
 317
 318       if (offset + compressedSize > end - data) {
 319         MOZ_ASSERT(false, "StartupCache file is corrupt.");
 320         return Err(NS_ERROR_UNEXPECTED);
 321       }
 322
 323       // Make sure offsets match what we'd expect based on script ordering and
 324       // size, as a basic sanity check.
 325       if (offset != currentOffset) {
 326         return Err(NS_ERROR_UNEXPECTED);
 327       }
 328       currentOffset += compressedSize;
 329
 330       // We could use mTable.putNew if we knew the file we're loading weren't
 331       // corrupt. However, we don't know that, so check if the key already
 332       // exists. If it does, we know the file must be corrupt.
 333       decltype(mTable)::AddPtr p = mTable.lookupForAdd(key);
 334       if (p) {
 335         return Err(NS_ERROR_UNEXPECTED);
 336       }
 337
 338       if (!mTable.add(
 339               p, key,
 340               StartupCacheEntry(offset, compressedSize, uncompressedSize))) {
 341         return Err(NS_ERROR_UNEXPECTED);
 342       }
 343     }
 344
 345     if (buf.error()) {
 346       return Err(NS_ERROR_UNEXPECTED);
 347     }
 348
 349     cleanup.release();
 350   }
 351
 352   MMAP_FAULT_HANDLER_CATCH(Err(NS_ERROR_UNEXPECTED))
 353
 354   return Ok();
 355 }
 356
 357 bool StartupCache::HasEntry(const char* id) {
 358   AUTO_PROFILER_LABEL("StartupCache::HasEntry", OTHER);
 359
 360   MOZ_ASSERT(NS_IsMainThread(), "Startup cache only available on main thread");
 361
 362   return mTable.has(nsDependentCString(id));
 363 }
 364
 365 nsresult StartupCache::GetBuffer(const char* id, const char** outbuf,
 366                                  uint32_t* length) {
 367   AUTO_PROFILER_LABEL("StartupCache::GetBuffer", OTHER);
 368
 369   NS_ASSERTION(NS_IsMainThread(),
 370                "Startup cache only available on main thread");
 371
 372   Telemetry::LABELS_STARTUP_CACHE_REQUESTS label =
 373       Telemetry::LABELS_STARTUP_CACHE_REQUESTS::Miss;
 374   auto telemetry =
 375       MakeScopeExit([&label] { Telemetry::AccumulateCategorical(label); });
 376
 377   decltype(mTable)::Ptr p = mTable.lookup(nsDependentCString(id));
 378   if (!p) {
 379     return NS_ERROR_NOT_AVAILABLE;
 380   }
 381
 382   auto& value = p->value();
 383   if (value.mData) {
 384     label = Telemetry::LABELS_STARTUP_CACHE_REQUESTS::HitMemory;
 385   } else {
 386     if (!mCacheData.initialized()) {
 387       return NS_ERROR_NOT_AVAILABLE;
 388     }
 389 #ifdef DEBUG
 390     // It should be impossible for a write to be pending here. This is because
 391     // we just checked mCacheData.initialized(), and this is reset before
 392     // writing to the cache. It's not re-initialized unless we call
 393     // LoadArchive(), either from Init() (which must have already happened) or
 394     // InvalidateCache(). InvalidateCache() locks the mutex, so a write can't be
 395     // happening. Really, we want to MOZ_ASSERT(!mTableLock.IsLocked()) here,
 396     // but there is no such method. So we hack around by attempting to gain the
 397     // lock. This should always succeed; if it fails, someone's broken the
 398     // assumptions.
 399     if (!mTableLock.TryLock()) {
 400       MOZ_ASSERT(false, "Could not gain mTableLock - should never happen!");
 401       return NS_ERROR_NOT_AVAILABLE;
 402     }
 403     mTableLock.Unlock();
 404 #endif
 405
 406     size_t totalRead = 0;
 407     size_t totalWritten = 0;
 408     Span<const char> compressed = Span(
 409         mCacheData.get<char>().get() + mCacheEntriesBaseOffset + value.mOffset,
 410         value.mCompressedSize);
 411     value.mData = MakeUnique<char[]>(value.mUncompressedSize);
 412     Span<char> uncompressed = Span(value.mData.get(), value.mUncompressedSize);
 413     MMAP_FAULT_HANDLER_BEGIN_BUFFER(uncompressed.Elements(),
 414                                     uncompressed.Length())
 415     bool finished = false;
 416     while (!finished) {
 417       auto result = mDecompressionContext->Decompress(
 418           uncompressed.From(totalWritten), compressed.From(totalRead));
 419       if (NS_WARN_IF(result.isErr())) {
 420         value.mData = nullptr;
 421         InvalidateCache();
 422         return NS_ERROR_FAILURE;
 423       }
 424       auto decompressionResult = result.unwrap();
 425       totalRead += decompressionResult.mSizeRead;
 426       totalWritten += decompressionResult.mSizeWritten;
 427       finished = decompressionResult.mFinished;
 428     }
 429
 430     MMAP_FAULT_HANDLER_CATCH(NS_ERROR_FAILURE)
 431
 432     label = Telemetry::LABELS_STARTUP_CACHE_REQUESTS::HitDisk;
 433   }
 434
 435   if (!value.mRequested) {
 436     value.mRequested = true;
 437     value.mRequestedOrder = ++mRequestedCount;
 438     MOZ_ASSERT(mRequestedCount <= mTable.count(),
 439                "Somehow we requested more StartupCache items than exist.");
 440     ResetStartupWriteTimerCheckingReadCount();
 441   }
 442
 443   // Track that something holds a reference into mTable, so we know to hold
 444   // onto it in case the cache is invalidated.
 445   mCurTableReferenced = true;
 446   *outbuf = value.mData.get();
 447   *length = value.mUncompressedSize;
 448   return NS_OK;
 449 }
 450
 451 // Makes a copy of the buffer, client retains ownership of inbuf.
 452 nsresult StartupCache::PutBuffer(const char* id, UniquePtr<char[]>&& inbuf,
 453                                  uint32_t len) {
 454   NS_ASSERTION(NS_IsMainThread(),
 455                "Startup cache only available on main thread");
 456   if (StartupCache::gShutdownInitiated) {
 457     return NS_ERROR_NOT_AVAILABLE;
 458   }
 459
 460   bool exists = mTable.has(nsDependentCString(id));
 461
 462   if (exists) {
 463     NS_WARNING("Existing entry in StartupCache.");
 464     // Double-caching is undesirable but not an error.
 465     return NS_OK;
 466   }
 467   // Try to gain the table write lock. If the background task to write the
 468   // cache is running, this will fail.
 469   if (!mTableLock.TryLock()) {
 470     return NS_ERROR_NOT_AVAILABLE;
 471   }
 472   auto lockGuard = MakeScopeExit([&] { mTableLock.Unlock(); });
 473
 474   // putNew returns false on alloc failure - in the very unlikely event we hit
 475   // that and aren't going to crash elsewhere, there's no reason we need to
 476   // crash here.
 477   if (mTable.putNew(nsCString(id), StartupCacheEntry(std::move(inbuf), len,
 478                                                      ++mRequestedCount))) {
 479     return ResetStartupWriteTimer();
 480   }
 481   MOZ_DIAGNOSTIC_ASSERT(mTable.count() < STARTUP_CACHE_MAX_CAPACITY,
 482                         "Too many StartupCache entries.");
 483   return NS_OK;
 484 }
 485
 486 size_t StartupCache::HeapSizeOfIncludingThis(
 487     mozilla::MallocSizeOf aMallocSizeOf) const {
 488   // This function could measure more members, but they haven't been found by
 489   // DMD to be significant.  They can be added later if necessary.
 490
 491   size_t n = aMallocSizeOf(this);
 492
 493   n += mTable.shallowSizeOfExcludingThis(aMallocSizeOf);
 494   for (auto iter = mTable.iter(); !iter.done(); iter.next()) {
 495     if (iter.get().value().mData) {
 496       n += aMallocSizeOf(iter.get().value().mData.get());
 497     }
 498     n += iter.get().key().SizeOfExcludingThisIfUnshared(aMallocSizeOf);
 499   }
 500
 501   return n;
 502 }
 503
 504 /**
 505  * WriteToDisk writes the cache out to disk. Callers of WriteToDisk need to call
 506  * WaitOnWriteComplete to make sure there isn't a write
 507  * happening on another thread
 508  */
 509 Result<Ok, nsresult> StartupCache::WriteToDisk() {
 510   mTableLock.AssertCurrentThreadOwns();
 511
 512   if (!mDirty || mWrittenOnce) {
 513     return Ok();
 514   }
 515
 516   if (!mFile) {
 517     return Err(NS_ERROR_UNEXPECTED);
 518   }
 519
 520   AutoFDClose fd;
 521   MOZ_TRY(mFile->OpenNSPRFileDesc(PR_WRONLY | PR_CREATE_FILE | PR_TRUNCATE,
 522                                   0644, &fd.rwget()));
 523
 524   nsTArray<std::pair<const nsCString*, StartupCacheEntry*>> entries;
 525   for (auto iter = mTable.iter(); !iter.done(); iter.next()) {
 526     if (iter.get().value().mRequested) {
 527       entries.AppendElement(
 528           std::make_pair(&iter.get().key(), &iter.get().value()));
 529     }
 530   }
 531
 532   if (entries.IsEmpty()) {
 533     return Ok();
 534   }
 535
 536   entries.Sort(StartupCacheEntry::Comparator());
 537   loader::OutputBuffer buf;
 538   for (auto& e : entries) {
 539     auto key = e.first;
 540     auto value = e.second;
 541     auto uncompressedSize = value->mUncompressedSize;
 542     // Set the mHeaderOffsetInFile so we can go back and edit the offset.
 543     value->mHeaderOffsetInFile = buf.cursor();
 544     // Write a 0 offset/compressed size as a placeholder until we get the real
 545     // offset after compressing.
 546     buf.codeUint32(0);
 547     buf.codeUint32(0);
 548     buf.codeUint32(uncompressedSize);
 549     buf.codeString(*key);
 550   }
 551
 552   uint8_t headerSize[4];
 553   LittleEndian::writeUint32(headerSize, buf.cursor());
 554
 555   MOZ_TRY(Write(fd, MAGIC, sizeof(MAGIC)));
 556   MOZ_TRY(Write(fd, headerSize, sizeof(headerSize)));
 557   size_t headerStart = sizeof(MAGIC) + sizeof(headerSize);
 558   size_t dataStart = headerStart + buf.cursor();
 559   MOZ_TRY(Seek(fd, dataStart));
 560
 561   size_t offset = 0;
 562
 563   const size_t chunkSize = 1024 * 16;
 564   LZ4FrameCompressionContext ctx(6,         /* aCompressionLevel */
 565                                  chunkSize, /* aReadBufLen */
 566                                  true,      /* aChecksum */
 567                                  true);     /* aStableSrc */
 568   size_t writeBufLen = ctx.GetRequiredWriteBufferLength();
 569   auto writeBuffer = MakeUnique<char[]>(writeBufLen);
 570   auto writeSpan = Span(writeBuffer.get(), writeBufLen);
 571
 572   for (auto& e : entries) {
 573     auto value = e.second;
 574     value->mOffset = offset;
 575     Span<const char> result;
 576     MOZ_TRY_VAR(result,
 577                 ctx.BeginCompressing(writeSpan).mapErr(MapLZ4ErrorToNsresult));
 578     MOZ_TRY(Write(fd, result.Elements(), result.Length()));
 579     offset += result.Length();
 580
 581     for (size_t i = 0; i < value->mUncompressedSize; i += chunkSize) {
 582       size_t size = std::min(chunkSize, value->mUncompressedSize - i);
 583       char* uncompressed = value->mData.get() + i;
 584       MOZ_TRY_VAR(result, ctx.ContinueCompressing(Span(uncompressed, size))
 585                               .mapErr(MapLZ4ErrorToNsresult));
 586       MOZ_TRY(Write(fd, result.Elements(), result.Length()));
 587       offset += result.Length();
 588     }
 589
 590     MOZ_TRY_VAR(result, ctx.EndCompressing().mapErr(MapLZ4ErrorToNsresult));
 591     MOZ_TRY(Write(fd, result.Elements(), result.Length()));
 592     offset += result.Length();
 593     value->mCompressedSize = offset - value->mOffset;
 594     MOZ_TRY(Seek(fd, dataStart + offset));
 595   }
 596
 597   for (auto& e : entries) {
 598     auto value = e.second;
 599     uint8_t* headerEntry = buf.Get() + value->mHeaderOffsetInFile;
 600     LittleEndian::writeUint32(headerEntry, value->mOffset);
 601     LittleEndian::writeUint32(headerEntry + sizeof(value->mOffset),
 602                               value->mCompressedSize);
 603   }
 604   MOZ_TRY(Seek(fd, headerStart));
 605   MOZ_TRY(Write(fd, buf.Get(), buf.cursor()));
 606
 607   mDirty = false;
 608   mWrittenOnce = true;
 609
 610   return Ok();
 611 }
 612
 613 void StartupCache::InvalidateCache(bool memoryOnly) {
 614   WaitOnPrefetchThread();
 615   // Ensure we're not writing using mTable...
 616   MutexAutoLock unlock(mTableLock);
 617
 618   mWrittenOnce = false;
 619   if (memoryOnly) {
 620     // This should only be called in tests.
 621     auto writeResult = WriteToDisk();
 622     if (NS_WARN_IF(writeResult.isErr())) {
 623       gIgnoreDiskCache = true;
 624       return;
 625     }
 626   }
 627   if (mCurTableReferenced) {
 628     // There should be no way for this assert to fail other than a user manually
 629     // sending startupcache-invalidate messages through the Browser Toolbox.
 630     MOZ_DIAGNOSTIC_ASSERT(xpc::IsInAutomation() || mOldTables.Length() < 10,
 631                           "Startup cache invalidated too many times.");
 632     mOldTables.AppendElement(std::move(mTable));
 633     mCurTableReferenced = false;
 634   } else {
 635     mTable.clear();
 636   }
 637   mRequestedCount = 0;
 638   if (!memoryOnly) {
 639     mCacheData.reset();
 640     nsresult rv = mFile->Remove(false);
 641     if (NS_FAILED(rv) && rv != NS_ERROR_FILE_TARGET_DOES_NOT_EXIST &&
 642         rv != NS_ERROR_FILE_NOT_FOUND) {
 643       gIgnoreDiskCache = true;
 644       return;
 645     }
 646   }
 647   gIgnoreDiskCache = false;
 648   auto result = LoadArchive();
 649   if (NS_WARN_IF(result.isErr())) {
 650     gIgnoreDiskCache = true;
 651   }
 652 }
 653
 654 void StartupCache::MaybeInitShutdownWrite() {
 655   if (mTimer) {
 656     mTimer->Cancel();
 657   }
 658   gShutdownInitiated = true;
 659
 660   MaybeWriteOffMainThread();
 661 }
 662
 663 void StartupCache::EnsureShutdownWriteComplete() {
 664   // If we've already written or there's nothing to write,
 665   // we don't need to do anything. This is the common case.
 666   if (mWrittenOnce || (mCacheData.initialized() && !ShouldCompactCache())) {
 667     return;
 668   }
 669   // Otherwise, ensure the write happens. The timer should have been cancelled
 670   // already in MaybeInitShutdownWrite.
 671   if (!mTableLock.TryLock()) {
 672     // Uh oh, we're writing away from the main thread. Wait to gain the lock,
 673     // to ensure the write completes.
 674     mTableLock.Lock();
 675   } else {
 676     // We got the lock. Keep the following in sync with
 677     // MaybeWriteOffMainThread:
 678     WaitOnPrefetchThread();
 679     mDirty = true;
 680     mCacheData.reset();
 681     // Most of this should be redundant given MaybeWriteOffMainThread should
 682     // have run before now.
 683
 684     auto writeResult = WriteToDisk();
 685     Unused << NS_WARN_IF(writeResult.isErr());
 686     // We've had the lock, and `WriteToDisk()` sets mWrittenOnce and mDirty
 687     // when done, and checks for them when starting, so we don't need to do
 688     // anything else.
 689   }
 690   mTableLock.Unlock();
 691 }
 692
 693 void StartupCache::IgnoreDiskCache() {
 694   gIgnoreDiskCache = true;
 695   if (gStartupCache) gStartupCache->InvalidateCache();
 696 }
 697
 698 void StartupCache::WaitOnPrefetchThread() {
 699   if (!mPrefetchThread || mPrefetchThread == PR_GetCurrentThread()) return;
 700
 701   PR_JoinThread(mPrefetchThread);
 702   mPrefetchThread = nullptr;
 703 }
 704
 705 void StartupCache::ThreadedPrefetch(void* aClosure) {
 706   AUTO_PROFILER_REGISTER_THREAD("StartupCache");
 707   NS_SetCurrentThreadName("StartupCache");
 708   mozilla::IOInterposer::RegisterCurrentThread();
 709   StartupCache* startupCacheObj = static_cast<StartupCache*>(aClosure);
 710   uint8_t* buf = startupCacheObj->mCacheData.get<uint8_t>().get();
 711   size_t size = startupCacheObj->mCacheData.size();
 712   MMAP_FAULT_HANDLER_BEGIN_BUFFER(buf, size)
 713   PrefetchMemory(buf, size);
 714   MMAP_FAULT_HANDLER_CATCH()
 715   mozilla::IOInterposer::UnregisterCurrentThread();
 716 }
 717
 718 bool StartupCache::ShouldCompactCache() {
 719   // If we've requested less than 4/5 of the startup cache, then we should
 720   // probably compact it down. This can happen quite easily after the first run,
 721   // which seems to request quite a few more things than subsequent runs.
 722   CheckedInt<uint32_t> threshold = CheckedInt<uint32_t>(mTable.count()) * 4 / 5;
 723   MOZ_RELEASE_ASSERT(threshold.isValid(), "Runaway StartupCache size");
 724   return mRequestedCount < threshold.value();
 725 }
 726
 727 /*
 728  * The write-thread is spawned on a timeout(which is reset with every write).
 729  * This can avoid a slow shutdown.
 730  */
 731 void StartupCache::WriteTimeout(nsITimer* aTimer, void* aClosure) {
 732   /*
 733    * It is safe to use the pointer passed in aClosure to reference the
 734    * StartupCache object because the timer's lifetime is tightly coupled to
 735    * the lifetime of the StartupCache object; this timer is canceled in the
 736    * StartupCache destructor, guaranteeing that this function runs if and only
 737    * if the StartupCache object is valid.
 738    */
 739   StartupCache* startupCacheObj = static_cast<StartupCache*>(aClosure);
 740   startupCacheObj->MaybeWriteOffMainThread();
 741 }
 742
 743 /*
 744  * See StartupCache::WriteTimeout above - this is just the non-static body.
 745  */
 746 void StartupCache::MaybeWriteOffMainThread() {
 747   if (mWrittenOnce) {
 748     return;
 749   }
 750
 751   if (mCacheData.initialized() && !ShouldCompactCache()) {
 752     return;
 753   }
 754
 755   // Keep this code in sync with EnsureShutdownWriteComplete.
 756   WaitOnPrefetchThread();
 757   mDirty = true;
 758   mCacheData.reset();
 759
 760   RefPtr<StartupCache> self = this;
 761   nsCOMPtr<nsIRunnable> runnable =
 762       NS_NewRunnableFunction("StartupCache::Write", [self]() mutable {
 763         MutexAutoLock unlock(self->mTableLock);
 764         auto result = self->WriteToDisk();
 765         Unused << NS_WARN_IF(result.isErr());
 766       });
 767   NS_DispatchBackgroundTask(runnable.forget(), NS_DISPATCH_EVENT_MAY_BLOCK);
 768 }
 769
 770 // We don't want to refcount StartupCache, so we'll just
 771 // hold a ref to this and pass it to observerService instead.
 772 NS_IMPL_ISUPPORTS(StartupCacheListener, nsIObserver)
 773
 774 nsresult StartupCacheListener::Observe(nsISupports* subject, const char* topic,
 775                                        const char16_t* data) {
 776   StartupCache* sc = StartupCache::GetSingleton();
 777   if (!sc) return NS_OK;
 778
 779   if (strcmp(topic, NS_XPCOM_SHUTDOWN_OBSERVER_ID) == 0) {
 780     // Do not leave the thread running past xpcom shutdown
 781     sc->WaitOnPrefetchThread();
 782     StartupCache::gShutdownInitiated = true;
 783     // Note that we don't do anything special for the background write
 784     // task; we expect the threadpool to finish running any tasks already
 785     // posted to it prior to shutdown. FastShutdown will call
 786     // EnsureShutdownWriteComplete() to ensure any pending writes happen
 787     // in that case.
 788   } else if (strcmp(topic, "startupcache-invalidate") == 0) {
 789     sc->InvalidateCache(data && nsCRT::strcmp(data, u"memoryOnly") == 0);
 790   }
 791   return NS_OK;
 792 }
 793
 794 nsresult StartupCache::GetDebugObjectOutputStream(
 795     nsIObjectOutputStream* aStream, nsIObjectOutputStream** aOutStream) {
 796   NS_ENSURE_ARG_POINTER(aStream);
 797 #ifdef DEBUG
 798   auto* stream = new StartupCacheDebugOutputStream(aStream, &mWriteObjectMap);
 799   NS_ADDREF(*aOutStream = stream);
 800 #else
 801   NS_ADDREF(*aOutStream = aStream);
 802 #endif
 803
 804   return NS_OK;
 805 }
 806
 807 nsresult StartupCache::ResetStartupWriteTimerCheckingReadCount() {
 808   nsresult rv = NS_OK;
 809   if (!mTimer)
 810     mTimer = NS_NewTimer();
 811   else
 812     rv = mTimer->Cancel();
 813   NS_ENSURE_SUCCESS(rv, rv);
 814   // Wait for the specified timeout, then write out the cache.
 815   mTimer->InitWithNamedFuncCallback(
 816       StartupCache::WriteTimeout, this, STARTUP_CACHE_WRITE_TIMEOUT * 1000,
 817       nsITimer::TYPE_ONE_SHOT, "StartupCache::WriteTimeout");
 818   return NS_OK;
 819 }
 820
 821 nsresult StartupCache::ResetStartupWriteTimer() {
 822   mDirty = true;
 823   nsresult rv = NS_OK;
 824   if (!mTimer)
 825     mTimer = NS_NewTimer();
 826   else
 827     rv = mTimer->Cancel();
 828   NS_ENSURE_SUCCESS(rv, rv);
 829   // Wait for the specified timeout, then write out the cache.
 830   mTimer->InitWithNamedFuncCallback(
 831       StartupCache::WriteTimeout, this, STARTUP_CACHE_WRITE_TIMEOUT * 1000,
 832       nsITimer::TYPE_ONE_SHOT, "StartupCache::WriteTimeout");
 833   return NS_OK;
 834 }
 835
 836 // Used only in tests:
 837 bool StartupCache::StartupWriteComplete() {
 838   // Need to have written to disk and not added new things since;
 839   return !mDirty && mWrittenOnce;
 840 }
 841
 842 // StartupCacheDebugOutputStream implementation
 843 #ifdef DEBUG
 844 NS_IMPL_ISUPPORTS(StartupCacheDebugOutputStream, nsIObjectOutputStream,
 845                   nsIBinaryOutputStream, nsIOutputStream)
 846
 847 bool StartupCacheDebugOutputStream::CheckReferences(nsISupports* aObject) {
 848   nsresult rv;
 849
 850   nsCOMPtr<nsIClassInfo> classInfo = do_QueryInterface(aObject);
 851   if (!classInfo) {
 852     NS_ERROR("aObject must implement nsIClassInfo");
 853     return false;
 854   }
 855
 856   uint32_t flags;
 857   rv = classInfo->GetFlags(&flags);
 858   NS_ENSURE_SUCCESS(rv, false);
 859   if (flags & nsIClassInfo::SINGLETON) return true;
 860
 861   bool inserted = mObjectMap->EnsureInserted(aObject);
 862   if (!inserted) {
 863     NS_ERROR(
 864         "non-singleton aObject is referenced multiple times in this"
 865         "serialization, we don't support that.");
 866   }
 867
 868   return inserted;
 869 }
 870
 871 // nsIObjectOutputStream implementation
 872 nsresult StartupCacheDebugOutputStream::WriteObject(nsISupports* aObject,
 873                                                     bool aIsStrongRef) {
 874   nsCOMPtr<nsISupports> rootObject(do_QueryInterface(aObject));
 875
 876   NS_ASSERTION(rootObject.get() == aObject,
 877                "bad call to WriteObject -- call WriteCompoundObject!");
 878   bool check = CheckReferences(aObject);
 879   NS_ENSURE_TRUE(check, NS_ERROR_FAILURE);
 880   return mBinaryStream->WriteObject(aObject, aIsStrongRef);
 881 }
 882
 883 nsresult StartupCacheDebugOutputStream::WriteSingleRefObject(
 884     nsISupports* aObject) {
 885   nsCOMPtr<nsISupports> rootObject(do_QueryInterface(aObject));
 886
 887   NS_ASSERTION(rootObject.get() == aObject,
 888                "bad call to WriteSingleRefObject -- call WriteCompoundObject!");
 889   bool check = CheckReferences(aObject);
 890   NS_ENSURE_TRUE(check, NS_ERROR_FAILURE);
 891   return mBinaryStream->WriteSingleRefObject(aObject);
 892 }
 893
 894 nsresult StartupCacheDebugOutputStream::WriteCompoundObject(
 895     nsISupports* aObject, const nsIID& aIID, bool aIsStrongRef) {
 896   nsCOMPtr<nsISupports> rootObject(do_QueryInterface(aObject));
 897
 898   nsCOMPtr<nsISupports> roundtrip;
 899   rootObject->QueryInterface(aIID, getter_AddRefs(roundtrip));
 900   NS_ASSERTION(roundtrip.get() == aObject,
 901                "bad aggregation or multiple inheritance detected by call to "
 902                "WriteCompoundObject!");
 903
 904   bool check = CheckReferences(aObject);
 905   NS_ENSURE_TRUE(check, NS_ERROR_FAILURE);
 906   return mBinaryStream->WriteCompoundObject(aObject, aIID, aIsStrongRef);
 907 }
 908
 909 nsresult StartupCacheDebugOutputStream::WriteID(nsID const& aID) {
 910   return mBinaryStream->WriteID(aID);
 911 }
 912
 913 char* StartupCacheDebugOutputStream::GetBuffer(uint32_t aLength,
 914                                                uint32_t aAlignMask) {
 915   return mBinaryStream->GetBuffer(aLength, aAlignMask);
 916 }
 917
 918 void StartupCacheDebugOutputStream::PutBuffer(char* aBuffer, uint32_t aLength) {
 919   mBinaryStream->PutBuffer(aBuffer, aLength);
 920 }
 921 #endif  // DEBUG
 922
 923 }  // namespace scache
 924 }  // namespace mozilla