Fix storing VkPipelineCache in HWUI's shader cache breaking persistence

See b/268205519 for the full analysis and description, but the summary
is mSavePending could previously get stuck in the true state, preventing
the shader cache from being written to disk for the rest of an app's
lifecycle.

This stuck state seems to occur on every app launch after the first,
unless any shaders are compiled in the first 4 seconds. This occurs
frequently, as the shaders initially required by an app in its first 4
seconds were likely compiled on its first launch and already available
in its persistent cache, thus not necessitating new compilation in that
initial window and triggering the failure case.

This seems to have serious jank implications, particularly for SysUI.

Note: time is now represented in milliseconds to make testing easier.

Fixes: 268205519

Bug: 268205519
Bug: 225211273

Test: `atest hwui_unit_tests:ShaderCacheTest` (before and after fix)

Change-Id: I102d75df4c61f9b3012dd5d956dc37435d5b5a84
diff --git a/libs/hwui/tests/unit/ShaderCacheTests.cpp b/libs/hwui/tests/unit/ShaderCacheTests.cpp
index 576e946..7bcd45c 100644
--- a/libs/hwui/tests/unit/ShaderCacheTests.cpp
+++ b/libs/hwui/tests/unit/ShaderCacheTests.cpp
@@ -14,6 +14,10 @@
  * limitations under the License.
  */
 
+#include <GrDirectContext.h>
+#include <Properties.h>
+#include <SkData.h>
+#include <SkRefCnt.h>
 #include <cutils/properties.h>
 #include <dirent.h>
 #include <errno.h>
@@ -22,11 +26,12 @@
 #include <stdlib.h>
 #include <sys/types.h>
 #include <utils/Log.h>
+
 #include <cstdint>
+
 #include "FileBlobCache.h"
 #include "pipeline/skia/ShaderCache.h"
-#include <SkData.h>
-#include <SkRefCnt.h>
+#include "tests/common/TestUtils.h"
 
 using namespace android::uirenderer::skiapipeline;
 
@@ -37,11 +42,38 @@
 class ShaderCacheTestUtils {
 public:
     /**
-     * "setSaveDelay" sets the time in seconds to wait before saving newly inserted cache entries.
-     * If set to 0, then deferred save is disabled.
+     * Hack to reset all member variables of the given cache to their default / initial values.
+     *
+     * WARNING: this must be kept up to date manually, since ShaderCache's parent disables just
+     * reassigning a new instance.
      */
-    static void setSaveDelay(ShaderCache& cache, unsigned int saveDelay) {
-        cache.mDeferredSaveDelay = saveDelay;
+    static void reinitializeAllFields(ShaderCache& cache) {
+        ShaderCache newCache = ShaderCache();
+        std::lock_guard<std::mutex> lock(cache.mMutex);
+        // By order of declaration
+        cache.mInitialized = newCache.mInitialized;
+        cache.mBlobCache.reset(nullptr);
+        cache.mFilename = newCache.mFilename;
+        cache.mIDHash.clear();
+        cache.mSavePending = newCache.mSavePending;
+        cache.mObservedBlobValueSize = newCache.mObservedBlobValueSize;
+        cache.mDeferredSaveDelayMs = newCache.mDeferredSaveDelayMs;
+        cache.mTryToStorePipelineCache = newCache.mTryToStorePipelineCache;
+        cache.mInStoreVkPipelineInProgress = newCache.mInStoreVkPipelineInProgress;
+        cache.mNewPipelineCacheSize = newCache.mNewPipelineCacheSize;
+        cache.mOldPipelineCacheSize = newCache.mOldPipelineCacheSize;
+        cache.mCacheDirty = newCache.mCacheDirty;
+        cache.mNumShadersCachedInRam = newCache.mNumShadersCachedInRam;
+    }
+
+    /**
+     * "setSaveDelayMs" sets the time in milliseconds to wait before saving newly inserted cache
+     * entries. If set to 0, then deferred save is disabled, and "saveToDiskLocked" must be called
+     * manually, as seen in the "terminate" testing helper function.
+     */
+    static void setSaveDelayMs(ShaderCache& cache, unsigned int saveDelayMs) {
+        std::lock_guard<std::mutex> lock(cache.mMutex);
+        cache.mDeferredSaveDelayMs = saveDelayMs;
     }
 
     /**
@@ -50,8 +82,9 @@
      */
     static void terminate(ShaderCache& cache, bool saveContent) {
         std::lock_guard<std::mutex> lock(cache.mMutex);
-        cache.mSavePending = saveContent;
-        cache.saveToDiskLocked();
+        if (saveContent) {
+            cache.saveToDiskLocked();
+        }
         cache.mBlobCache = NULL;
     }
 
@@ -62,6 +95,38 @@
     static bool validateCache(ShaderCache& cache, std::vector<T> hash) {
         return cache.validateCache(hash.data(), hash.size() * sizeof(T));
     }
+
+    /**
+     * Waits until cache::mSavePending is false, checking every 0.1 ms *while the mutex is free*.
+     *
+     * Fails if there was no save pending, or if the cache was already being written to disk, or if
+     * timeoutMs is exceeded.
+     *
+     * Note: timeoutMs only guards against mSavePending getting stuck like in b/268205519, and
+     * cannot protect against mutex-based deadlock. Reaching timeoutMs implies something is broken,
+     * so setting it to a sufficiently large value will not delay execution in the happy state.
+     */
+    static void waitForPendingSave(ShaderCache& cache, const int timeoutMs = 50) {
+        {
+            std::lock_guard<std::mutex> lock(cache.mMutex);
+            ASSERT_TRUE(cache.mSavePending);
+        }
+        bool saving = true;
+        float elapsedMilliseconds = 0;
+        while (saving) {
+            if (elapsedMilliseconds >= timeoutMs) {
+                FAIL() << "Timed out after waiting " << timeoutMs << " ms for a pending save";
+            }
+            // This small (0.1 ms) delay is to avoid working too much while waiting for
+            // deferredSaveThread to take the mutex and start the disk write.
+            const int delayMicroseconds = 100;
+            usleep(delayMicroseconds);
+            elapsedMilliseconds += (float)delayMicroseconds / 1000;
+
+            std::lock_guard<std::mutex> lock(cache.mMutex);
+            saving = cache.mSavePending;
+        }
+    }
 };
 
 } /* namespace skiapipeline */
@@ -83,6 +148,18 @@
     return false;
 }
 
+/**
+ * Attempts to delete the given file, and asserts that either:
+ * 1. Deletion was successful, OR
+ * 2. The file did not exist.
+ *
+ * Tip: wrap calls to this in ASSERT_NO_FATAL_FAILURE(x) if a test should exit early if this fails.
+ */
+void deleteFileAssertSuccess(const std::string& filePath) {
+    int deleteResult = remove(filePath.c_str());
+    ASSERT_TRUE(0 == deleteResult || ENOENT == errno);
+}
+
 inline bool checkShader(const sk_sp<SkData>& shader1, const sk_sp<SkData>& shader2) {
     return nullptr != shader1 && nullptr != shader2 && shader1->size() == shader2->size() &&
            0 == memcmp(shader1->data(), shader2->data(), shader1->size());
@@ -93,6 +170,10 @@
     return checkShader(shader, shader2);
 }
 
+inline bool checkShader(const sk_sp<SkData>& shader, const std::string& program) {
+    return checkShader(shader, program.c_str());
+}
+
 template <typename T>
 bool checkShader(const sk_sp<SkData>& shader, std::vector<T>& program) {
     sk_sp<SkData> shader2 = SkData::MakeWithCopy(program.data(), program.size() * sizeof(T));
@@ -103,6 +184,10 @@
     shader = SkData::MakeWithCString(program);
 }
 
+void setShader(sk_sp<SkData>& shader, const std::string& program) {
+    setShader(shader, program.c_str());
+}
+
 template <typename T>
 void setShader(sk_sp<SkData>& shader, std::vector<T>& buffer) {
     shader = SkData::MakeWithCopy(buffer.data(), buffer.size() * sizeof(T));
@@ -126,13 +211,13 @@
     std::string cacheFile2 = getExternalStorageFolder() + "/shaderCacheTest2";
 
     // remove any test files from previous test run
-    int deleteFile = remove(cacheFile1.c_str());
-    ASSERT_TRUE(0 == deleteFile || ENOENT == errno);
+    ASSERT_NO_FATAL_FAILURE(deleteFileAssertSuccess(cacheFile1));
+    ASSERT_NO_FATAL_FAILURE(deleteFileAssertSuccess(cacheFile2));
     std::srand(0);
 
     // read the cache from a file that does not exist
     ShaderCache::get().setFilename(cacheFile1.c_str());
-    ShaderCacheTestUtils::setSaveDelay(ShaderCache::get(), 0);  // disable deferred save
+    ShaderCacheTestUtils::setSaveDelayMs(ShaderCache::get(), 0);  // disable deferred save
     ShaderCache::get().initShaderDiskCache();
 
     // read a key - should not be found since the cache is empty
@@ -186,7 +271,8 @@
     ASSERT_TRUE(checkShader(outVS2, dataBuffer));
 
     ShaderCacheTestUtils::terminate(ShaderCache::get(), false);
-    remove(cacheFile1.c_str());
+    ASSERT_NO_FATAL_FAILURE(deleteFileAssertSuccess(cacheFile1));
+    ASSERT_NO_FATAL_FAILURE(deleteFileAssertSuccess(cacheFile2));
 }
 
 TEST(ShaderCacheTest, testCacheValidation) {
@@ -198,13 +284,13 @@
     std::string cacheFile2 = getExternalStorageFolder() + "/shaderCacheTest2";
 
     // remove any test files from previous test run
-    int deleteFile = remove(cacheFile1.c_str());
-    ASSERT_TRUE(0 == deleteFile || ENOENT == errno);
+    ASSERT_NO_FATAL_FAILURE(deleteFileAssertSuccess(cacheFile1));
+    ASSERT_NO_FATAL_FAILURE(deleteFileAssertSuccess(cacheFile2));
     std::srand(0);
 
     // generate identity and read the cache from a file that does not exist
     ShaderCache::get().setFilename(cacheFile1.c_str());
-    ShaderCacheTestUtils::setSaveDelay(ShaderCache::get(), 0);  // disable deferred save
+    ShaderCacheTestUtils::setSaveDelayMs(ShaderCache::get(), 0);  // disable deferred save
     std::vector<uint8_t> identity(1024);
     genRandomData(identity);
     ShaderCache::get().initShaderDiskCache(
@@ -278,7 +364,81 @@
     }
 
     ShaderCacheTestUtils::terminate(ShaderCache::get(), false);
-    remove(cacheFile1.c_str());
+    ASSERT_NO_FATAL_FAILURE(deleteFileAssertSuccess(cacheFile1));
+    ASSERT_NO_FATAL_FAILURE(deleteFileAssertSuccess(cacheFile2));
+}
+
+using namespace android::uirenderer;
+RENDERTHREAD_SKIA_PIPELINE_TEST(ShaderCacheTest, testOnVkFrameFlushed) {
+    if (Properties::getRenderPipelineType() != RenderPipelineType::SkiaVulkan) {
+        // RENDERTHREAD_SKIA_PIPELINE_TEST declares both SkiaVK and SkiaGL variants.
+        GTEST_SKIP() << "This test is only applicable to RenderPipelineType::SkiaVulkan";
+    }
+    if (!folderExist(getExternalStorageFolder())) {
+        // Don't run the test if external storage folder is not available
+        return;
+    }
+    std::string cacheFile = getExternalStorageFolder() + "/shaderCacheTest";
+    GrDirectContext* grContext = renderThread.getGrContext();
+
+    // Remove any test files from previous test run
+    ASSERT_NO_FATAL_FAILURE(deleteFileAssertSuccess(cacheFile));
+
+    // The first iteration of this loop is to save an initial VkPipelineCache data blob to disk,
+    // which sets up the second iteration for a common scenario of comparing a "new" VkPipelineCache
+    // blob passed to "store" against the same blob that's already in the persistent cache from a
+    // previous launch. "reinitializeAllFields" is critical to emulate each iteration being as close
+    // to the state of a freshly launched app as possible, as the initial values of member variables
+    // like mInStoreVkPipelineInProgress and mOldPipelineCacheSize are critical to catch issues
+    // such as b/268205519
+    for (int flushIteration = 1; flushIteration <= 2; flushIteration++) {
+        SCOPED_TRACE("Frame flush iteration " + std::to_string(flushIteration));
+        // Reset *all* in-memory data and reload the cache from disk.
+        ShaderCacheTestUtils::reinitializeAllFields(ShaderCache::get());
+        ShaderCacheTestUtils::setSaveDelayMs(ShaderCache::get(), 10);  // Delay must be > 0 to save.
+        ShaderCache::get().setFilename(cacheFile.c_str());
+        ShaderCache::get().initShaderDiskCache();
+
+        // 1st iteration: store pipeline data to be read back on a subsequent "boot" of the "app".
+        // 2nd iteration: ensure that an initial frame flush (without storing any shaders) given the
+        // same pipeline data that's already on disk doesn't break the cache.
+        ShaderCache::get().onVkFrameFlushed(grContext);
+        ASSERT_NO_FATAL_FAILURE(ShaderCacheTestUtils::waitForPendingSave(ShaderCache::get()));
+    }
+
+    constexpr char shader1[] = "sassas";
+    constexpr char shader2[] = "someVS";
+    constexpr int numIterations = 3;
+    // Also do n iterations of separate "store some shaders then flush the frame" pairs to just
+    // double-check the cache also doesn't get stuck from that use case.
+    for (int saveIteration = 1; saveIteration <= numIterations; saveIteration++) {
+        SCOPED_TRACE("Shader save iteration " + std::to_string(saveIteration));
+        // Write twice to the in-memory cache, which should start a deferred save with both queued.
+        sk_sp<SkData> inVS;
+        setShader(inVS, shader1 + std::to_string(saveIteration));
+        ShaderCache::get().store(GrProgramDescTest(100), *inVS.get(), SkString());
+        setShader(inVS, shader2 + std::to_string(saveIteration));
+        ShaderCache::get().store(GrProgramDescTest(432), *inVS.get(), SkString());
+
+        // Simulate flush to also save latest pipeline info.
+        ShaderCache::get().onVkFrameFlushed(grContext);
+        ASSERT_NO_FATAL_FAILURE(ShaderCacheTestUtils::waitForPendingSave(ShaderCache::get()));
+    }
+
+    // Reload from disk to ensure saving succeeded.
+    ShaderCacheTestUtils::terminate(ShaderCache::get(), false);
+    ShaderCache::get().initShaderDiskCache();
+
+    // Read twice, ensure equal to last store.
+    sk_sp<SkData> outVS;
+    ASSERT_NE((outVS = ShaderCache::get().load(GrProgramDescTest(100))), sk_sp<SkData>());
+    ASSERT_TRUE(checkShader(outVS, shader1 + std::to_string(numIterations)));
+    ASSERT_NE((outVS = ShaderCache::get().load(GrProgramDescTest(432))), sk_sp<SkData>());
+    ASSERT_TRUE(checkShader(outVS, shader2 + std::to_string(numIterations)));
+
+    // Clean up.
+    ShaderCacheTestUtils::terminate(ShaderCache::get(), false);
+    ASSERT_NO_FATAL_FAILURE(deleteFileAssertSuccess(cacheFile));
 }
 
 }  // namespace