FileStream-based implementation for updating cache index file (#2485)

* Draft FileStream-based implementation for updating cache file

* File streams fully integrated into shader cache code paths; Tests will not run unless file system is on disk as file streams do not play nicely with in-memory

* Brought old code back as fallback path, but tests need to ensure previous is freed first

* Testing structure updated, beginning cleanup work

* All tests working

* Cleanup changes

* Removed an extra tab at the end of a line

* Cleanup change

* Undo externals change

* Removed redundant logic for OS vs memory file system handling of the shader cache; Removed extra helper function left over from old cache implementation

* Reverted performance change to generate contents hashes when modules are being loaded as this code path is not always followed; Contents hashing now uses a combination of hashing and checking the last modified time for all file dependencies, only reading in and hashing the contents of all files if the last modified hash does not match

* Added handling to Module::updateContentsBasedHash for file dependencies which are not from a physical source file on disk; Added test for above

Co-authored-by: Lucy Chen <lucchen@nvidia.com>
Co-authored-by: Yong He <yonghe@outlook.com>

1 files changed, 195 insertions, 68 deletions

diff --git a/tools/gfx/persistent-shader-cache.cpp b/tools/gfx/persistent-shader-cache.cpp
index c0fd3b44a..54c46f9dc 100644
--- a/tools/gfx/persistent-shader-cache.cpp
+++ b/tools/gfx/persistent-shader-cache.cpp
@@ -6,9 +6,15 @@
 #include "../../source/core/slang-string-util.h"
 #include "../../source/core/slang-file-system.h"
 
+#include "../../source/core/slang-char-util.h"
+
+#include <chrono>
+
 namespace gfx
 {
 
+using namespace std::chrono;
+
 PersistentShaderCache::PersistentShaderCache(const IDevice::ShaderCacheDesc& inDesc)
 {
     desc = inDesc;
@@ -35,52 +41,97 @@ PersistentShaderCache::PersistentShaderCache(const IDevice::ShaderCacheDesc& inD
     loadCacheFromFile();
 }
 
+PersistentShaderCache::~PersistentShaderCache()
+{
+    if (isMemoryFileSystem)
+    {
+        saveCacheToMemory();
+    }
+}
+
 // Load a previous cache index saved to disk. If not found, create a new cache index
 // and save it to disk as filename.
 void PersistentShaderCache::loadCacheFromFile()
 {
-    ComPtr<ISlangBlob> indexBlob;
-    if (SLANG_FAILED(desc.shaderCacheFileSystem->loadFile(desc.cacheFilename, indexBlob.writeRef())))
+    // We will need to combine the filename with the cache path in order to have the correct
+    // file path for initializing the stream. This needs to be done separately because there
+    // is no guarantee that the underlying file system is mutable.
+    String filePath;
+    if (mutableShaderCacheFileSystem)
     {
-        // Cache index not found, so we'll create and save a new one.
-        if (mutableShaderCacheFileSystem)
+        ComPtr<ISlangBlob> fullPath;
+        if (SLANG_FAILED(mutableShaderCacheFileSystem->getPath(PathKind::OperatingSystem, desc.cacheFilename, fullPath.writeRef())))
         {
-            mutableShaderCacheFileSystem->saveFile(desc.cacheFilename, nullptr, 0);
+            // If we fail to obtain a physical file path, then this must be a MemoryFileSystem. In this case, file streams
+            // will not work as they require the file to be on disk, so we will rely on a fall back implementation.
+            isMemoryFileSystem = true;
+            loadCacheFromMemory();
             return;
         }
-        // Cache index not found and we can't save a new one due to the file system being immutable.
-        return;
+        filePath = String((char*)fullPath->getBufferPointer());
+    }
+    else
+    {
+        filePath = Path::combine(String(desc.shaderCachePath), String(desc.cacheFilename));
     }
 
-    auto indexString = UnownedStringSlice((char*)indexBlob->getBufferPointer());
-
-    List<UnownedStringSlice> lines;
-    StringUtil::calcLines(indexString, lines);
-    for (auto line : lines)
+    if (SLANG_FAILED(indexStream.init(filePath, FileMode::Open, FileAccess::ReadWrite, FileShare::ReadWrite)))
     {
-        List<UnownedStringSlice> digests;
-        StringUtil::split(line, ' ', digests); // This will return our two hashes as two elements in digests, unless we've reached the end.
-        if (digests.getCount() != 2)
-            continue;
-        auto dependencyDigest = DigestUtil::fromString(digests[0]);
-        auto contentsDigest = DigestUtil::fromString(digests[1]);
+        // If we failed to open a stream to the file, then the file does not yet exist on disk.
+        // We will create the index file if our underlying file system is mutable.
+        if (mutableShaderCacheFileSystem)
+        {
+            indexStream.init(filePath, FileMode::Create, FileAccess::ReadWrite, FileShare::ReadWrite);
+        }
+        return;
+    }
+    else
+    {
+        const auto start = indexStream.getPosition();
+        indexStream.seek(SeekOrigin::End, 0);
+        const auto end = indexStream.getPosition();
+        indexStream.seek(SeekOrigin::Start, 0);
+        const Index numEntries = (Index)(end - start) / sizeof(ShaderCacheEntry);
 
-        ShaderCacheEntry entry = { dependencyDigest, contentsDigest };
-        auto entryNode = entries.AddLast(entry);
-        keyToEntry.Add(dependencyDigest, entryNode);
+        if (desc.entryCountLimit > 0 && numEntries > desc.entryCountLimit)
+        {
+            // If the size limit for the current cache is smaller than the cache that produced the file we're trying to
+            // load, re-create the entire file.
+            //
+            // FileStream does not currently have any methods for truncating an existing file, so in this case, our cache
+            // index would no longer accurately reflect the state of our cache due to the extra now-garbage lines present.
+            // While this has no impact on cache operation, it could be problematic for debugging purposes, etc.
+            indexStream.close();
+            indexStream.init(filePath, FileMode::Create, FileAccess::ReadWrite, FileShare::ReadWrite);
+            return;
+        }
+        else
+        {
+            // The cache index is not guaranteed to be ordered by most recent access, so we need a temporary list to store
+            // all the entries in order to sort them before filling in our linked list.
+            List<ShaderCacheEntry> tempEntries;
+            tempEntries.setCount(numEntries);
+            size_t bytesRead;
+            indexStream.read(tempEntries.getBuffer(), sizeof(ShaderCacheEntry) * numEntries, bytesRead);
 
-        // If there are more entries present in the cache file than the entry count limit allows,
-        // ignore the rest. Since we output the lines in the cache file in order of most to least
-        // recently used, the cache's behavior will still be correct.
-        if (desc.entryCountLimit > 0 && entries.Count() == desc.entryCountLimit)
-            break;
+            // We will need to sort tempEntries by last accessed time before we can add entries to our linked list.
+            tempEntries.quickSort(tempEntries.getBuffer(), 0, tempEntries.getCount() - 1, [](ShaderCacheEntry a, ShaderCacheEntry b) { return a.lastAccessedTime > b.lastAccessedTime; });
+            for (auto& entry : tempEntries)
+            {
+                // If we reach this point, then the current cache is at least the same size in entries as the cache
+                // that produced the index we're reading in, so we don't need to check if we're exceeding capacity.
+                auto entryIndexNode = orderedEntries.AddLast(entries.getCount());
+                entries.add(entry);
+                keyToEntry.Add(entry.dependencyBasedDigest, entryIndexNode);
+            }
+        }   
     }
 }
 
-LinkedNode<ShaderCacheEntry>* PersistentShaderCache::findEntry(const slang::Digest& key, ISlangBlob** outCompiledCode)
+ShaderCacheEntry* PersistentShaderCache::findEntry(const slang::Digest& key, ISlangBlob** outCompiledCode)
 {
-    LinkedNode<ShaderCacheEntry>* entryNode;
-    if (!keyToEntry.TryGetValue(key, entryNode))
+    LinkedNode<Index>* entryIndexNode;
+    if (!keyToEntry.TryGetValue(key, entryIndexNode))
     {
         // The key was not found in the cache, so we return nullptr.
         *outCompiledCode = nullptr;
@@ -90,19 +141,24 @@ LinkedNode<ShaderCacheEntry>* PersistentShaderCache::findEntry(const slang::Dige
     // If the key is found, load the stored contents from disk. We then move the corresponding
     // entry to the front of the linked list and update the cache file on disk
     desc.shaderCacheFileSystem->loadFile(DigestUtil::toString(key).getBuffer(), outCompiledCode);
-    if (entries.FirstNode() != entryNode)
+    auto index = entryIndexNode->Value;
+    entries[index].lastAccessedTime = (double)high_resolution_clock::now().time_since_epoch().count();
+    if (orderedEntries.FirstNode() != entryIndexNode)
     {
-        entries.RemoveFromList(entryNode);
-        entries.AddFirst(entryNode);
-        if (mutableShaderCacheFileSystem)
+        orderedEntries.RemoveFromList(entryIndexNode);
+        orderedEntries.AddFirst(entryIndexNode);
+        if (mutableShaderCacheFileSystem && !isMemoryFileSystem)
         {
-            saveCacheToFile();
+            auto offset = index * sizeof(ShaderCacheEntry);
+            indexStream.seek(SeekOrigin::Start, offset + 2 * sizeof(slang::Digest));
+            indexStream.write(&entries[index].lastAccessedTime, sizeof(double));
+            indexStream.flush();
         }
     }
-    return entryNode;
+    return &entries[index];
 }
 
-void PersistentShaderCache::addEntry(const slang::Digest& dependencyDigest, const slang::Digest& astDigest, ISlangBlob* compiledCode)
+void PersistentShaderCache::addEntry(const slang::Digest& dependencyDigest, const slang::Digest& contentsDigest, ISlangBlob* compiledCode)
 {
     if (!mutableShaderCacheFileSystem)
     {
@@ -113,25 +169,45 @@ void PersistentShaderCache::addEntry(const slang::Digest& dependencyDigest, cons
     // Check that we do not exceed the cache's size limit by adding another entry. If so,
     // remove the least recently used entry first.
     //
-    // In theory, this could loop infinitely since deleteLRUEntry() immediately returns if
-    // mutableShaderCacheFileSystem is not set. However, this situation is functionally impossible
-    // because we check immediately before this as well.
-    while (desc.entryCountLimit > 0 && entries.Count() >= desc.entryCountLimit)
+    // In theory, the cache could be more than just one entry over the entry count limit.
+    // However, this is impossible in practice because we fully re-create the entry list
+    // and cache index file if the size of the current cache is smaller than the cache
+    // that generated the index file we loaded. In any case, the initial number of entries
+    // in the cache will always be fewer than the size limit and this check will be hit
+    // on the first entry added that exceeds the cache's size.
+    Index index = entries.getCount();
+    if (desc.entryCountLimit > 0 && orderedEntries.Count() >= desc.entryCountLimit)
     {
-        deleteLRUEntry();
+        index = deleteLRUEntry();
     }
 
-    ShaderCacheEntry entry = { dependencyDigest, astDigest };
-    auto entryNode = entries.AddFirst(entry);
+    auto lastAccessedTime = (double)high_resolution_clock::now().time_since_epoch().count();
+
+    ShaderCacheEntry entry = { dependencyDigest, contentsDigest, lastAccessedTime };
+    auto entryNode = orderedEntries.AddFirst(index);
+    if (index == entries.getCount())
+    {
+        // No entries were removed, so we can tack this entry on at the end.
+        entries.add(entry);
+    }
+    else
+    {
+        // An entry was deleted, so we overwrite that slot with the new entry.
+        entries[index] = entry;
+    }
     keyToEntry.Add(dependencyDigest, entryNode);
 
     mutableShaderCacheFileSystem->saveFileBlob(DigestUtil::toString(dependencyDigest).getBuffer(), compiledCode);
 
-    saveCacheToFile();
+    if (!isMemoryFileSystem)
+    {
+        indexStream.seek(SeekOrigin::End, 0);
+        indexStream.write(&entry, sizeof(ShaderCacheEntry));
+        indexStream.flush();
+    }
 }
 
 void PersistentShaderCache::updateEntry(
-    LinkedNode<ShaderCacheEntry>* entryNode,
     const slang::Digest& dependencyDigest,
     const slang::Digest& contentsDigest,
     ISlangBlob* updatedCode)
@@ -143,48 +219,99 @@ void PersistentShaderCache::updateEntry(
         return;
     }
 
-    entryNode->Value.contentsBasedDigest = contentsDigest;
+    // Unlike in addEntry(), we only update the contents digest here because the last accessed time will have already
+    // been updated while finding the entry.
+    auto entryIndexNode = *keyToEntry.TryGetValue(dependencyDigest);
+    auto index = entryIndexNode->Value;
+    entries[index].contentsBasedDigest = contentsDigest;
     mutableShaderCacheFileSystem->saveFileBlob(DigestUtil::toString(dependencyDigest).getBuffer(), updatedCode);
 
-    saveCacheToFile();
+    if (!isMemoryFileSystem)
+    {
+        auto offset = index * sizeof(ShaderCacheEntry);
+        indexStream.seek(SeekOrigin::Start, offset + sizeof(slang::Digest));
+        indexStream.write(&contentsDigest, sizeof(slang::Digest));
+        indexStream.flush();
+    }
 }
 
-void PersistentShaderCache::saveCacheToFile()
+Index PersistentShaderCache::deleteLRUEntry()
 {
     if (!mutableShaderCacheFileSystem)
     {
-        // Cannot save the index to disk if the underlying file system isn't mutable.
-        return;
+        // This is here as a safety precaution but should never be hit as
+        // addEntry() and its memory-based equivalent are the only functions
+        // that should call this.
+        return -1;
     }
 
-    StringBuilder indexSb;
-    for (auto& entry : entries)
-    {
-        indexSb << entry.dependencyBasedDigest;
-        indexSb << " ";
-        indexSb << entry.contentsBasedDigest;
-        indexSb << "\n";
-    }
+    auto lruEntry = orderedEntries.LastNode();
+    auto index = lruEntry->Value;
+    auto shaderKey = entries[index].dependencyBasedDigest;
 
-    mutableShaderCacheFileSystem->saveFile(desc.cacheFilename, indexSb.getBuffer(), indexSb.getLength());
+    keyToEntry.Remove(shaderKey);
+    mutableShaderCacheFileSystem->remove(DigestUtil::toString(shaderKey).getBuffer());
+
+    orderedEntries.Delete(lruEntry);
+    return index;
 }
 
-void PersistentShaderCache::deleteLRUEntry()
+// An in-memory file system cannot utilize file streaming to update the index file in place.
+// Consequently, the cache index file is updated once on exit and is guaranteed to maintain the
+// correct order of entries from most to least recently used. However, any kind of interruption
+// in program execution that results in the cache destructor not being called will result in an
+// inaccurate cache index.
+// 
+// These currently assume that the underlying file system must be a MemoryFileSystem as this is the
+// only in-memory file system that currently exists in Slang, which is guaranteed to be mutable.
+// Mutability checks will need to be added if this changes in the future.
+void PersistentShaderCache::loadCacheFromMemory()
 {
-    if (!mutableShaderCacheFileSystem)
+    ComPtr<ISlangBlob> indexBlob;
+    if (SLANG_FAILED(mutableShaderCacheFileSystem->loadFile(desc.cacheFilename, indexBlob.writeRef())))
     {
-        // This is here as a safety precaution but should never be hit as
-        // addEntry() is the only function that should call this.
+        mutableShaderCacheFileSystem->saveFile(desc.cacheFilename, nullptr, 0);
         return;
     }
 
-    auto lruEntry = entries.LastNode();
-    auto shaderKey = lruEntry->Value.dependencyBasedDigest;
+    auto indexString = UnownedStringSlice((char*)indexBlob->getBufferPointer());
 
-    keyToEntry.Remove(shaderKey);
-    mutableShaderCacheFileSystem->remove(DigestUtil::toString(shaderKey).getBuffer());
+    List<UnownedStringSlice> lines;
+    StringUtil::calcLines(indexString, lines);
+    for (auto line : lines)
+    {
+        List<UnownedStringSlice> entryFields;
+        StringUtil::split(line, ' ', entryFields);
+        if (entryFields.getCount() != 2)
+            continue;
+
+        ShaderCacheEntry entry;
+        entry.dependencyBasedDigest = DigestUtil::fromString(entryFields[0]);
+        entry.contentsBasedDigest = DigestUtil::fromString(entryFields[1]);
+        entry.lastAccessedTime = 0;
+
+        auto entryNode = orderedEntries.AddLast(entries.getCount());
+        entries.add(entry);
+        keyToEntry.Add(entry.dependencyBasedDigest, entryNode);
+
+        if (desc.entryCountLimit > 0 && orderedEntries.Count() == desc.entryCountLimit)
+            break;
+    }
+}
 
-    entries.Delete(lruEntry);
+void PersistentShaderCache::saveCacheToMemory()
+{
+    StringBuilder indexSb;
+    for (auto& entryIndex : orderedEntries)
+    {
+        auto entry = entries[entryIndex];
+        indexSb << entry.dependencyBasedDigest;
+        indexSb << " ";
+        indexSb << entry.contentsBasedDigest;
+        indexSb << "\n";
+    }
+
+    mutableShaderCacheFileSystem->saveFile(desc.cacheFilename, indexSb.getBuffer(), indexSb.getLength());
 }
 
 }