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