Merge dynamic manifest updates to Github (#7)

This change introduces dynamic manifest updates to asset streaming.

Asset streaming describes the directory to be streamed in a manifest, which is a proto definition of all content metadata. This information is sufficient to answer `stat` and `readdir` calls in the FUSE layer without additional round-trips to the workstation.

When a directory is streamed for the first time, the corresponding manifest is created in two steps:
1. The directory is traversed recursively and the inode information of all contained files and directories is written to the manifest.
2. The content of all identified files is processed to generate each file's chunk list. This list is part of the definition of a file in the manifest.
  * The chunk boundaries are identified using our implementation of the FastCDC algorithm.
  * The hash of each chunk is calculated using the BLAKE3 hash function.
  * The length and hash of each chunk is appended to the file's chunk list.

Prior to this change, when the user mounted a workstation directory on a client, the asset streaming server pushed an intermediate manifest to the gamelet as soon as step 1 was completed. At this point, the FUSE client started serving the virtual file system and was ready to answer `stat` and `readdir` calls. In case the FUSE client received any call that required file contents, such as `read`, it would block the caller until the server completed step 2 above and pushed the final manifest to the client. This works well for large directories (> 100GB) with a reasonable number of files (< 100k). But when dealing with millions of tiny files, creating the full manifest can take several minutes.

With this change, we introduce dynamic manifest updates. When the FUSE layer receives an `open` or `readdir` request for a file or directory that is incomplete, it sends an RPC to the workstation about what information is missing from the manifest. The workstation identifies the corresponding file chunker or directory scanner tasks and moves them to the front of the queue. As soon as the task is completed, the workstation pushes an updated intermediate manifest to the client which now includes the information to serve the FUSE request. The queued FUSE request is resumed and returns the result to the caller.

While this does not reduce the required time to build the final manifest, it splits up the work into smaller tasks. This allows us to interrupt the current work and prioritize those tasks which are required to handle an incoming request from the client. While this still takes a round-trip to the workstation plus the processing time for the task, an updated manifest is received within a few seconds, which is much better than blocking for several minutes. 

This latency is only visible when serving data while the manifest is still being created. The situation improves as the manifest creation on the workstation progresses. As soon as the final manifest is pushed, all metadata can be served directly without having to wait for pending tasks.

asset_stream_manager/asset_stream_server.cc

@@ -26,11 +26,11 @@ AssetStreamServer::AssetStreamServer(std::string src_dir,
 std::unique_ptr<AssetStreamServer> AssetStreamServer::Create(
     AssetStreamServerType type, std::string src_dir,
     DataStoreReader* data_store_reader, FileChunkMap* file_chunks,
-    ContentSentHandler content_sent) {
+    ContentSentHandler content_sent, PrioritizeAssetsHandler prio_assets) {
   switch (type) {
     case AssetStreamServerType::kGrpc:
-      return std::make_unique<GrpcAssetStreamServer>(src_dir, data_store_reader,
-                                                     file_chunks, content_sent);
+      return std::make_unique<GrpcAssetStreamServer>(
+          src_dir, data_store_reader, file_chunks, content_sent, prio_assets);
     case AssetStreamServerType::kTest:
       return std::make_unique<TestingAssetStreamServer>(
           src_dir, data_store_reader, file_chunks);
@@ -38,4 +38,5 @@ std::unique_ptr<AssetStreamServer> AssetStreamServer::Create(
   assert(false);
   return nullptr;
 }
+
 }  // namespace cdc_ft

asset_stream_manager/asset_stream_server.h

@@ -33,6 +33,11 @@ namespace cdc_ft {
 using ContentSentHandler = std::function<void(
     size_t byte_count, size_t chunk_count, std::string instance_id)>;
 
+// Handles requests to prioritize the given list of assets while updating the
+// manifest. |rel_paths| relative Unix paths of assets to prioritize.
+using PrioritizeAssetsHandler =
+    std::function<void(std::vector<std::string> rel_paths)>;
+
 class DataStoreReader;
 class FileChunkMap;
 
@@ -49,7 +54,7 @@ class AssetStreamServer {
   static std::unique_ptr<AssetStreamServer> Create(
       AssetStreamServerType type, std::string src_dir,
       DataStoreReader* data_store_reader, FileChunkMap* file_chunks,
-      ContentSentHandler content_sent);
+      ContentSentHandler content_sent, PrioritizeAssetsHandler prio_assets);
 
   AssetStreamServer(const AssetStreamServer& other) = delete;
   AssetStreamServer& operator=(const AssetStreamServer& other) = delete;

asset_stream_manager/grpc_asset_stream_server.cc

@@ -40,6 +40,8 @@ using GetManifestIdResponse = proto::GetManifestIdResponse;
 using AckManifestIdReceivedRequest = proto::AckManifestIdReceivedRequest;
 using AckManifestIdReceivedResponse = proto::AckManifestIdReceivedResponse;
 using ConfigStreamService = proto::ConfigStreamService;
+using ProcessAssetsRequest = proto::ProcessAssetsRequest;
+using ProcessAssetsResponse = proto::ProcessAssetsResponse;
 
 }  // namespace
 
@@ -64,6 +66,8 @@ class AssetStreamServiceImpl final : public AssetStreamService::Service {
     std::string rel_path;
     uint64_t offset;
     size_t size;
+    std::string instance_id = instance_ids_->Get(context->peer());
+
     for (const ContentIdProto& id : request->id()) {
       uint32_t uint32_size;
       if (file_chunks_->Lookup(id, &rel_path, &offset, &uint32_size)) {
@@ -77,7 +81,6 @@ class AssetStreamServiceImpl final : public AssetStreamService::Service {
         RETURN_GRPC_IF_ERROR(
             ReadFromDataStore(id, response->add_data(), &size));
       }
-      std::string instance_id = instance_ids_->Get(context->peer());
       if (content_sent_ != nullptr) {
         content_sent_(size, 1, instance_id);
       }
@@ -144,8 +147,9 @@ class AssetStreamServiceImpl final : public AssetStreamService::Service {
 
 class ConfigStreamServiceImpl final : public ConfigStreamService::Service {
  public:
-  explicit ConfigStreamServiceImpl(InstanceIdMap* instance_ids)
-      : instance_ids_(instance_ids) {}
+  ConfigStreamServiceImpl(InstanceIdMap* instance_ids,
+                          PrioritizeAssetsHandler prio_handler)
+      : instance_ids_(instance_ids), prio_handler_(std::move(prio_handler)) {}
   ~ConfigStreamServiceImpl() { Shutdown(); }
 
   grpc::Status GetManifestId(
@@ -183,6 +187,20 @@ class ConfigStreamServiceImpl final : public ConfigStreamService::Service {
     return grpc::Status::OK;
   }
 
+  grpc::Status ProcessAssets(grpc::ServerContext* context,
+                             const ProcessAssetsRequest* request,
+                             ProcessAssetsResponse* response) override {
+    if (!prio_handler_) return grpc::Status::OK;
+
+    std::vector<std::string> rel_paths;
+    rel_paths.reserve(request->relative_paths().size());
+    for (const std::string& rel_path : request->relative_paths()) {
+      rel_paths.push_back(rel_path);
+    }
+    prio_handler_(std::move(rel_paths));
+    return grpc::Status::OK;
+  }
+
   void SetManifestId(const ContentIdProto& id) ABSL_LOCKS_EXCLUDED(mutex_) {
     LOG_INFO("Updating manifest id '%s' in configuration service",
              ContentId::ToHexString(id));
@@ -219,6 +237,10 @@ class ConfigStreamServiceImpl final : public ConfigStreamService::Service {
     return id_;
   }
 
+  void SetPrioritizeAssetsHandler(PrioritizeAssetsHandler handler) {
+    prio_handler_ = handler;
+  }
+
  private:
   // Returns false if the update process was cancelled.
   bool WaitForUpdate(ContentIdProto& local_id) ABSL_LOCKS_EXCLUDED(mutex_) {
@@ -234,23 +256,24 @@ class ConfigStreamServiceImpl final : public ConfigStreamService::Service {
   mutable absl::Mutex mutex_;
   ContentIdProto id_ ABSL_GUARDED_BY(mutex_);
   bool running_ ABSL_GUARDED_BY(mutex_) = true;
-  InstanceIdMap* instance_ids_;
+  InstanceIdMap* instance_ids_ = nullptr;
+  PrioritizeAssetsHandler prio_handler_;
 
   // Maps instance ids to the last acknowledged manifest id.
   using AckedManifestIdsMap = std::unordered_map<std::string, ContentIdProto>;
   AckedManifestIdsMap acked_manifest_ids_ ABSL_GUARDED_BY(mutex_);
 };
 
-GrpcAssetStreamServer::GrpcAssetStreamServer(std::string src_dir,
-                                             DataStoreReader* data_store_reader,
-                                             FileChunkMap* file_chunks,
-                                             ContentSentHandler content_sent)
+GrpcAssetStreamServer::GrpcAssetStreamServer(
+    std::string src_dir, DataStoreReader* data_store_reader,
+    FileChunkMap* file_chunks, ContentSentHandler content_sent,
+    PrioritizeAssetsHandler prio_assets)
     : AssetStreamServer(src_dir, data_store_reader, file_chunks),
       asset_stream_service_(std::make_unique<AssetStreamServiceImpl>(
           std::move(src_dir), data_store_reader, file_chunks, &instance_ids_,
           content_sent)),
-      config_stream_service_(
-          std::make_unique<ConfigStreamServiceImpl>(&instance_ids_)) {}
+      config_stream_service_(std::make_unique<ConfigStreamServiceImpl>(
+          &instance_ids_, std::move(prio_assets))) {}
 
 GrpcAssetStreamServer::~GrpcAssetStreamServer() = default;
 

asset_stream_manager/grpc_asset_stream_server.h

@@ -40,7 +40,8 @@ class GrpcAssetStreamServer : public AssetStreamServer {
   // Creates a new asset streaming gRpc server.
   GrpcAssetStreamServer(std::string src_dir, DataStoreReader* data_store_reader,
                         FileChunkMap* file_chunks,
-                        ContentSentHandler content_sent);
+                        ContentSentHandler content_sent,
+                        PrioritizeAssetsHandler prio_assets);
 
   ~GrpcAssetStreamServer();
 

asset_stream_manager/local_assets_stream_manager_service_impl.h

@@ -57,15 +57,13 @@ class LocalAssetsStreamManagerServiceImpl final
   // if it exists.
   grpc::Status StartSession(grpc::ServerContext* context,
                             const StartSessionRequest* request,
-                            StartSessionResponse* response) override
-      ABSL_LOCKS_EXCLUDED(sessions_mutex_);
+                            StartSessionResponse* response) override;
 
   // Stops the streaming session to the instance with id
   // |request->gamelet_id()|. Returns a NotFound error if no session exists.
   grpc::Status StopSession(grpc::ServerContext* context,
                            const StopSessionRequest* request,
-                           StopSessionResponse* response) override
-      ABSL_LOCKS_EXCLUDED(sessions_mutex_);
+                           StopSessionResponse* response) override;
 
  private:
   // Convert StartSessionRequest enum to metrics enum.

asset_stream_manager/multi_session.cc

@@ -96,10 +96,25 @@ MultiSessionRunner::MultiSessionRunner(
 absl::Status MultiSessionRunner::Initialize(int port,
                                             AssetStreamServerType type,
                                             ContentSentHandler content_sent) {
+  // Create the manifest updater.
+  UpdaterConfig cfg;
+  cfg.num_threads = num_updater_threads_;
+  cfg.src_dir = src_dir_;
+
+  assert(!manifest_updater_);
+  manifest_updater_ =
+      std::make_unique<ManifestUpdater>(data_store_, std::move(cfg));
+
+  // Let the manifest updater handle requests to prioritize certain assets.
+  PrioritizeAssetsHandler prio_assets =
+      std::bind(&ManifestUpdater::AddPriorityAssets, manifest_updater_.get(),
+                std::placeholders::_1);
+
   // Start the server.
   assert(!server_);
   server_ = AssetStreamServer::Create(type, src_dir_, data_store_,
-                                      &file_chunks_, content_sent);
+                                      &file_chunks_, std::move(content_sent),
+                                      std::move(prio_assets));
   assert(server_);
   RETURN_IF_ERROR(server_->Start(port),
                   "Failed to start asset stream server for '%s'", src_dir_);
@@ -162,12 +177,6 @@ ContentIdProto MultiSessionRunner::ManifestId() const {
 }
 
 void MultiSessionRunner::Run() {
-  // Create the manifest updater.
-  UpdaterConfig cfg;
-  cfg.num_threads = num_updater_threads_;
-  cfg.src_dir = src_dir_;
-  ManifestUpdater manifest_updater(data_store_, std::move(cfg));
-
   // Set up file watcher.
   // The streamed path should be a directory and exist at the beginning.
   FileWatcherWin watcher(src_dir_);
@@ -179,28 +188,23 @@ void MultiSessionRunner::Run() {
     return;
   }
 
-  // Push an intermediate manifest containing the full directory structure, but
-  // potentially missing chunks. The purpose is that the FUSE can immediately
-  // show the structure and inode stats. FUSE will block on file reads that
-  // cannot be served due to missing chunks until the manifest is ready.
-  auto push_intermediate_manifest = [this](const ContentIdProto& manifest_id) {
+  // Push the intermediate manifest(s) and the final version with this handler.
+  auto push_handler = [this](const ContentIdProto& manifest_id) {
     SetManifest(manifest_id);
   };
 
   // Bring the manifest up to date.
   LOG_INFO("Updating manifest for '%s'...", src_dir_);
   Stopwatch sw;
-  status =
-      manifest_updater.UpdateAll(&file_chunks_, push_intermediate_manifest);
-  RecordManifestUpdate(manifest_updater, sw.Elapsed(),
+  status = manifest_updater_->UpdateAll(&file_chunks_, push_handler);
+  RecordManifestUpdate(*manifest_updater_, sw.Elapsed(),
                        metrics::UpdateTrigger::kInitUpdateAll, status);
   if (!status.ok()) {
     SetStatus(
         WrapStatus(status, "Failed to update manifest for '%s'", src_dir_));
     return;
   }
-  RecordMultiSessionStart(manifest_updater);
-  SetManifest(manifest_updater.ManifestId());
+  RecordMultiSessionStart(*manifest_updater_);
   LOG_INFO("Manifest for '%s' updated in %0.3f seconds", src_dir_,
            sw.ElapsedSeconds());
 
@@ -256,30 +260,26 @@ void MultiSessionRunner::Run() {
           src_dir_);
       modified_files.clear();
       sw.Reset();
-      status = manifest_updater.UpdateAll(&file_chunks_);
-      RecordManifestUpdate(manifest_updater, sw.Elapsed(),
+      status = manifest_updater_->UpdateAll(&file_chunks_, push_handler);
+      RecordManifestUpdate(*manifest_updater_, sw.Elapsed(),
                            metrics::UpdateTrigger::kRunningUpdateAll, status);
       if (!status.ok()) {
         LOG_WARNING(
             "Updating manifest for '%s' after re-creating directory failed: "
             "'%s'",
             src_dir_, status.ToString());
-        SetManifest(manifest_updater.DefaultManifestId());
-      } else {
-        SetManifest(manifest_updater.ManifestId());
+        SetManifest(manifest_updater_->DefaultManifestId());
       }
     } else if (!modified_files.empty()) {
       ManifestUpdater::OperationList ops = GetFileOperations(modified_files);
       sw.Reset();
-      status = manifest_updater.Update(&ops, &file_chunks_);
-      RecordManifestUpdate(manifest_updater, sw.Elapsed(),
+      status = manifest_updater_->Update(&ops, &file_chunks_, push_handler);
+      RecordManifestUpdate(*manifest_updater_, sw.Elapsed(),
                            metrics::UpdateTrigger::kRegularUpdate, status);
       if (!status.ok()) {
         LOG_WARNING("Updating manifest for '%s' failed: %s", src_dir_,
                     status.ToString());
-        SetManifest(manifest_updater.DefaultManifestId());
-      } else {
-        SetManifest(manifest_updater.ManifestId());
+        SetManifest(manifest_updater_->DefaultManifestId());
       }
     }
 
@@ -482,15 +482,16 @@ absl::Status MultiSession::Shutdown() {
     sessions_.erase(instance_id);
   }
 
+  absl::Status status;
   if (runner_) {
-    RETURN_IF_ERROR(runner_->Shutdown());
+    status = runner_->Shutdown();
   }
 
   if (heartbeat_watcher_.joinable()) {
     heartbeat_watcher_.join();
   }
 
-  return absl::OkStatus();
+  return status;
 }
 
 absl::Status MultiSession::Status() {
@@ -528,7 +529,7 @@ absl::Status MultiSession::StartSession(const std::string& instance_id,
   RETURN_IF_ERROR(session->Start(local_asset_stream_port_,
                                  kAssetStreamPortFirst, kAssetStreamPortLast));
 
-  // Wait for the FUSE to receive the intermediate manifest.
+  // Wait for the FUSE to receive the first intermediate manifest.
   RETURN_IF_ERROR(runner_->WaitForManifestAck(instance_id, absl::Seconds(5)));
 
   sessions_[instance_id] = std::move(session);

asset_stream_manager/multi_session.h

@@ -112,6 +112,7 @@ class MultiSessionRunner {
   const uint32_t num_updater_threads_;
   const ManifestUpdatedCb manifest_updated_cb_;
   std::unique_ptr<AssetStreamServer> server_;
+  std::unique_ptr<ManifestUpdater> manifest_updater_;
 
   // Modifications (shutdown, file changes).
   absl::Mutex mutex_;

asset_stream_manager/multi_session_test.cc

@@ -65,6 +65,21 @@ class MetricsServiceForTest : public MultiSessionMetricsRecorder {
     metrics_records_.push_back(MetricsRecord(std::move(event), code));
   }
 
+  // Waits until |num_events| events of type |type| have been recorded, or until
+  // the function times out. Returns true if the condition was met and false if
+  // in case of a timeout.
+  bool WaitForEvents(metrics::EventType type, int num_events = 1,
+                     absl::Duration timeout = absl::Seconds(1)) {
+    absl::MutexLock lock(&mutex_);
+    auto cond = [this, type, num_events]() {
+      return std::count_if(metrics_records_.begin(), metrics_records_.end(),
+                           [type](const MetricsRecord& mr) {
+                             return mr.code == type;
+                           }) >= num_events;
+    };
+    return mutex_.AwaitWithTimeout(absl::Condition(&cond), timeout);
+  }
+
   std::vector<MetricsRecord> GetEventsAndClear(metrics::EventType type)
       ABSL_LOCKS_EXCLUDED(mutex_) {
     std::vector<MetricsRecord> events;
@@ -172,8 +187,16 @@ class MultiSessionTest : public ManifestTestBase {
       metrics::ManifestUpdateData* data =
           events[i].evt.as_manager_data->manifest_update_data.get();
       EXPECT_LT(data->local_duration_ms, 60000ull);
-      manifests[i].local_duration_ms = data->local_duration_ms;
-      EXPECT_EQ(*data, manifests[i]);
+      EXPECT_EQ(data->status, manifests[i].status);
+      EXPECT_EQ(data->total_assets_added_or_updated,
+                manifests[i].total_assets_added_or_updated);
+      EXPECT_EQ(data->total_assets_deleted, manifests[i].total_assets_deleted);
+      EXPECT_EQ(data->total_chunks, manifests[i].total_chunks);
+      EXPECT_EQ(data->total_files_added_or_updated,
+                manifests[i].total_files_added_or_updated);
+      EXPECT_EQ(data->total_processed_bytes,
+                manifests[i].total_processed_bytes);
+      EXPECT_EQ(data->trigger, manifests[i].trigger);
     }
   }
 
@@ -268,13 +291,13 @@ TEST_F(MultiSessionTest, GetCachePath_DoesNotSplitUtfCodePoints) {
 TEST_F(MultiSessionTest, MultiSessionRunnerOnEmpty) {
   cfg_.src_dir = test_dir_path_;
   MultiSessionRunner runner(cfg_.src_dir, &data_store_, &process_factory_,
-                            false /*enable_stats*/, kTimeout, kNumThreads,
+                            /*enable_stats=*/false, kTimeout, kNumThreads,
                             metrics_service_,
                             [this]() { OnManifestUpdated(); });
   EXPECT_OK(runner.Initialize(kPort, AssetStreamServerType::kTest));
-  EXPECT_OK(runner.WaitForManifestAck(kInstance, kTimeout));
-  // The first update is always the empty manifest, wait for the second one.
-  ASSERT_TRUE(WaitForManifestUpdated(2));
+  EXPECT_TRUE(WaitForManifestUpdated(2));
+  ASSERT_TRUE(
+      metrics_service_->WaitForEvents(metrics::EventType::kMultiSessionStart));
   ASSERT_NO_FATAL_FAILURE(ExpectManifestEquals({}, runner.ManifestId()));
   CheckMultiSessionStartRecorded(0, 0, 0);
   CheckManifestUpdateRecorded(std::vector<metrics::ManifestUpdateData>{
@@ -290,13 +313,13 @@ TEST_F(MultiSessionTest, MultiSessionRunnerNonEmptySucceeds) {
   // Contains a.txt, subdir/b.txt, subdir/c.txt, subdir/d.txt.
   cfg_.src_dir = path::Join(base_dir_, "non_empty");
   MultiSessionRunner runner(cfg_.src_dir, &data_store_, &process_factory_,
-                            false /*enable_stats*/, kTimeout, kNumThreads,
+                            /*enable_stats=*/false, kTimeout, kNumThreads,
                             metrics_service_,
                             [this]() { OnManifestUpdated(); });
   EXPECT_OK(runner.Initialize(kPort, AssetStreamServerType::kTest));
-  EXPECT_OK(runner.WaitForManifestAck(kInstance, kTimeout));
-  // The first update is always the empty manifest, wait for the second one.
-  ASSERT_TRUE(WaitForManifestUpdated(2));
+  EXPECT_TRUE(WaitForManifestUpdated(2));
+  ASSERT_TRUE(
+      metrics_service_->WaitForEvents(metrics::EventType::kMultiSessionStart));
   CheckMultiSessionStartRecorded(46, 4, 4);
   ASSERT_NO_FATAL_FAILURE(ExpectManifestEquals(
       {"a.txt", "subdir", "subdir/b.txt", "subdir/c.txt", "subdir/d.txt"},
@@ -309,31 +332,50 @@ TEST_F(MultiSessionTest, MultiSessionRunnerNonEmptySucceeds) {
 TEST_F(MultiSessionTest, MultiSessionRunnerAddFileSucceeds) {
   cfg_.src_dir = test_dir_path_;
   MultiSessionRunner runner(cfg_.src_dir, &data_store_, &process_factory_,
-                            false /*enable_stats*/, kTimeout, kNumThreads,
+                            /*enable_stats=*/false, kTimeout, kNumThreads,
                             metrics_service_,
                             [this]() { OnManifestUpdated(); });
-  EXPECT_OK(runner.Initialize(kPort, AssetStreamServerType::kTest));
-  EXPECT_OK(runner.WaitForManifestAck(kInstance, kTimeout));
-  // The first update is always the empty manifest, wait for the second one.
-  ASSERT_TRUE(WaitForManifestUpdated(2));
-  ASSERT_OK(runner.Status());
-  CheckMultiSessionStartRecorded(0, 0, 0);
-  ASSERT_NO_FATAL_FAILURE(ExpectManifestEquals({}, runner.ManifestId()));
-  CheckManifestUpdateRecorded(std::vector<metrics::ManifestUpdateData>{
-      GetManifestUpdateData(metrics::UpdateTrigger::kInitUpdateAll,
-                            absl::StatusCode::kOk, 0, 0, 0, 0, 0, 0)});
 
-  const std::string file_path = path::Join(test_dir_path_, "file.txt");
-  EXPECT_OK(path::WriteFile(file_path, kData, kDataSize));
-  // 1 file was added = incremented exp_num_manifest_updates.
-  ASSERT_TRUE(WaitForManifestUpdated(3));
-  ASSERT_NO_FATAL_FAILURE(
-      ExpectManifestEquals({"file.txt"}, runner.ManifestId()));
-  CheckMultiSessionStartNotRecorded();
-  CheckManifestUpdateRecorded(std::vector<metrics::ManifestUpdateData>{
-      GetManifestUpdateData(metrics::UpdateTrigger::kRegularUpdate,
-                            absl::StatusCode::kOk, 1, 0, 1, 1, 0, kDataSize)});
+  {
+    SCOPED_TRACE("Initialize.");
 
+    EXPECT_OK(runner.Initialize(kPort, AssetStreamServerType::kTest));
+    // 1 file was added, 1 intermediate + 1 final manifest is pushed.
+    EXPECT_TRUE(WaitForManifestUpdated(2));
+    EXPECT_OK(runner.WaitForManifestAck(kInstance, kTimeout));
+    EXPECT_TRUE(metrics_service_->WaitForEvents(
+        metrics::EventType::kMultiSessionStart));
+    ASSERT_OK(runner.Status());
+  }
+
+  {
+    SCOPED_TRACE("Created base manifest for the test directory.");
+
+    CheckMultiSessionStartRecorded(0, 0, 0);
+    ASSERT_NO_FATAL_FAILURE(ExpectManifestEquals({}, runner.ManifestId()));
+    CheckManifestUpdateRecorded(std::vector<metrics::ManifestUpdateData>{
+        GetManifestUpdateData(metrics::UpdateTrigger::kInitUpdateAll,
+                              absl::StatusCode::kOk, 0, 0, 0, 0, 0, 0)});
+  }
+
+  {
+    SCOPED_TRACE("Added file.txt.");
+
+    uint32_t prev_updates = num_manifest_updates_;
+    const std::string file_path = path::Join(test_dir_path_, "file.txt");
+    EXPECT_OK(path::WriteFile(file_path, kData, kDataSize));
+    // 1 file was added, 1 intermediate + 1 final manifest is pushed.
+    EXPECT_TRUE(WaitForManifestUpdated(prev_updates + 2));
+    EXPECT_TRUE(
+        metrics_service_->WaitForEvents(metrics::EventType::kManifestUpdated));
+    ASSERT_NO_FATAL_FAILURE(
+        ExpectManifestEquals({"file.txt"}, runner.ManifestId()));
+    CheckMultiSessionStartNotRecorded();
+    CheckManifestUpdateRecorded(
+        std::vector<metrics::ManifestUpdateData>{GetManifestUpdateData(
+            metrics::UpdateTrigger::kRegularUpdate, absl::StatusCode::kOk, 1, 0,
+            1, 1, 0, kDataSize)});
+  }
   EXPECT_OK(runner.Status());
   EXPECT_OK(runner.Shutdown());
 }
@@ -343,7 +385,7 @@ TEST_F(MultiSessionTest, MultiSessionRunnerAddFileSucceeds) {
 TEST_F(MultiSessionTest, MultiSessionRunnerNoDirFails) {
   cfg_.src_dir = path::Join(base_dir_, "non_existing");
   MultiSessionRunner runner(cfg_.src_dir, &data_store_, &process_factory_,
-                            false /*enable_stats*/, kTimeout, kNumThreads,
+                            /*enable_stats=*/false, kTimeout, kNumThreads,
                             metrics_service_,
                             [this]() { OnManifestUpdated(); });
   EXPECT_OK(runner.Initialize(kPort, AssetStreamServerType::kTest));
@@ -365,16 +407,16 @@ TEST_F(MultiSessionTest, MultiSessionRunnerDirRecreatedSucceeds) {
                             kDataSize));
 
   MultiSessionRunner runner(cfg_.src_dir, &data_store_, &process_factory_,
-                            false /*enable_stats*/, kTimeout, kNumThreads,
+                            /*enable_stats=*/false, kTimeout, kNumThreads,
                             metrics_service_,
                             [this]() { OnManifestUpdated(); });
   EXPECT_OK(runner.Initialize(kPort, AssetStreamServerType::kTest));
 
   {
     SCOPED_TRACE("Originally, only the streamed directory contains file.txt.");
-    EXPECT_OK(runner.WaitForManifestAck(kInstance, kTimeout));
-    // The first update is always the empty manifest, wait for the second one.
-    ASSERT_TRUE(WaitForManifestUpdated(2));
+    EXPECT_TRUE(WaitForManifestUpdated(2));
+    ASSERT_TRUE(metrics_service_->WaitForEvents(
+        metrics::EventType::kMultiSessionStart));
     CheckMultiSessionStartRecorded((uint64_t)kDataSize, 1, 1);
     ASSERT_NO_FATAL_FAILURE(
         ExpectManifestEquals({"file.txt"}, runner.ManifestId()));
@@ -387,8 +429,9 @@ TEST_F(MultiSessionTest, MultiSessionRunnerDirRecreatedSucceeds) {
   {
     SCOPED_TRACE(
         "Remove the streamed directory, the manifest should become empty.");
+    uint32_t prev_updates = num_manifest_updates_;
     EXPECT_OK(path::RemoveDirRec(test_dir_path_));
-    ASSERT_TRUE(WaitForManifestUpdated(3));
+    ASSERT_TRUE(WaitForManifestUpdated(prev_updates + 1));
     ASSERT_NO_FATAL_FAILURE(ExpectManifestEquals({}, runner.ManifestId()));
     CheckManifestUpdateRecorded(
         std::vector<metrics::ManifestUpdateData>{GetManifestUpdateData(
@@ -400,9 +443,13 @@ TEST_F(MultiSessionTest, MultiSessionRunnerDirRecreatedSucceeds) {
     SCOPED_TRACE(
         "Create the watched directory -> an empty manifest should be "
         "streamed.");
+    uint32_t prev_updates = num_manifest_updates_;
     EXPECT_OK(path::CreateDirRec(test_dir_path_));
-    EXPECT_TRUE(WaitForManifestUpdated(4));
+    // The first update is always the empty manifest, wait for the second one.
+    EXPECT_TRUE(WaitForManifestUpdated(prev_updates + 2));
     ASSERT_NO_FATAL_FAILURE(ExpectManifestEquals({}, runner.ManifestId()));
+    EXPECT_TRUE(
+        metrics_service_->WaitForEvents(metrics::EventType::kManifestUpdated));
     CheckManifestUpdateRecorded(std::vector<metrics::ManifestUpdateData>{
         GetManifestUpdateData(metrics::UpdateTrigger::kRunningUpdateAll,
                               absl::StatusCode::kOk, 0, 0, 0, 0, 0, 0)});
@@ -410,11 +457,16 @@ TEST_F(MultiSessionTest, MultiSessionRunnerDirRecreatedSucceeds) {
 
   {
     SCOPED_TRACE("Create 'new_file.txt' -> new manifest should be created.");
+    uint32_t prev_updates = num_manifest_updates_;
     EXPECT_OK(path::WriteFile(path::Join(test_dir_path_, "new_file.txt"), kData,
                               kDataSize));
-    ASSERT_TRUE(WaitForManifestUpdated(5));
+    // The first update doesn't have the chunks for new_file.txt, wait for the
+    // second one.
+    ASSERT_TRUE(WaitForManifestUpdated(prev_updates + 2));
     ASSERT_NO_FATAL_FAILURE(
         ExpectManifestEquals({"new_file.txt"}, runner.ManifestId()));
+    EXPECT_TRUE(
+        metrics_service_->WaitForEvents(metrics::EventType::kManifestUpdated));
     CheckManifestUpdateRecorded(
         std::vector<metrics::ManifestUpdateData>{GetManifestUpdateData(
             metrics::UpdateTrigger::kRegularUpdate, absl::StatusCode::kOk, 1, 0,
@@ -432,11 +484,11 @@ TEST_F(MultiSessionTest, MultiSessionRunnerFileAsStreamedDirFails) {
   EXPECT_OK(path::WriteFile(cfg_.src_dir, kData, kDataSize));
 
   MultiSessionRunner runner(cfg_.src_dir, &data_store_, &process_factory_,
-                            false /*enable_stats*/, kTimeout, kNumThreads,
+                            /*enable_stats=*/false, kTimeout, kNumThreads,
                             metrics_service_,
                             [this]() { OnManifestUpdated(); });
   EXPECT_OK(runner.Initialize(kPort, AssetStreamServerType::kTest));
-  ASSERT_FALSE(WaitForManifestUpdated(1, absl::Milliseconds(10)));
+  ASSERT_FALSE(WaitForManifestUpdated(1, absl::Milliseconds(100)));
   CheckMultiSessionStartNotRecorded();
   CheckManifestUpdateRecorded(std::vector<metrics::ManifestUpdateData>{});
   EXPECT_NOT_OK(runner.Shutdown());
@@ -452,33 +504,49 @@ TEST_F(MultiSessionTest,
   EXPECT_OK(path::CreateDirRec(cfg_.src_dir));
 
   MultiSessionRunner runner(cfg_.src_dir, &data_store_, &process_factory_,
-                            false /*enable_stats*/, kTimeout, kNumThreads,
+                            /*enable_stats=*/false, kTimeout, kNumThreads,
                             metrics_service_,
                             [this]() { OnManifestUpdated(); });
-  EXPECT_OK(runner.Initialize(kPort, AssetStreamServerType::kTest));
-  ASSERT_TRUE(WaitForManifestUpdated(2));
-  CheckMultiSessionStartRecorded(0, 0, 0);
-  CheckManifestUpdateRecorded(std::vector<metrics::ManifestUpdateData>{
-      GetManifestUpdateData(metrics::UpdateTrigger::kInitUpdateAll,
-                            absl::StatusCode::kOk, 0, 0, 0, 0, 0, 0)});
-  ASSERT_NO_FATAL_FAILURE(ExpectManifestEquals({}, runner.ManifestId()));
+  {
+    SCOPED_TRACE("Initialize manifest in test directory.");
 
-  // Remove the streamed directory, the manifest should become empty.
-  EXPECT_OK(path::RemoveDirRec(cfg_.src_dir));
-  ASSERT_TRUE(WaitForManifestUpdated(3));
-  ASSERT_NO_FATAL_FAILURE(ExpectManifestEquals({}, runner.ManifestId()));
-  CheckManifestUpdateRecorded(std::vector<metrics::ManifestUpdateData>{
-      GetManifestUpdateData(metrics::UpdateTrigger::kRunningUpdateAll,
-                            absl::StatusCode::kNotFound, 0, 0, 0, 0, 0, 0)});
+    EXPECT_OK(runner.Initialize(kPort, AssetStreamServerType::kTest));
+    ASSERT_TRUE(WaitForManifestUpdated(2));
+    ASSERT_TRUE(metrics_service_->WaitForEvents(
+        metrics::EventType::kMultiSessionStart));
+    CheckMultiSessionStartRecorded(0, 0, 0);
+    CheckManifestUpdateRecorded(std::vector<metrics::ManifestUpdateData>{
+        GetManifestUpdateData(metrics::UpdateTrigger::kInitUpdateAll,
+                              absl::StatusCode::kOk, 0, 0, 0, 0, 0, 0)});
+    ASSERT_NO_FATAL_FAILURE(ExpectManifestEquals({}, runner.ManifestId()));
+  }
 
-  EXPECT_OK(path::WriteFile(cfg_.src_dir, kData, kDataSize));
-  EXPECT_TRUE(WaitForManifestUpdated(4));
-  ASSERT_NO_FATAL_FAILURE(ExpectManifestEquals({}, runner.ManifestId()));
-  CheckManifestUpdateRecorded(
-      std::vector<metrics::ManifestUpdateData>{GetManifestUpdateData(
-          metrics::UpdateTrigger::kRunningUpdateAll,
-          absl::StatusCode::kFailedPrecondition, 0, 0, 0, 0, 0, 0)});
-  CheckMultiSessionStartNotRecorded();
+  {
+    SCOPED_TRACE("Remove the streamed directory, the manifest becomes empty.");
+
+    uint32_t prev_updates = num_manifest_updates_;
+    EXPECT_OK(path::RemoveDirRec(cfg_.src_dir));
+    ASSERT_TRUE(WaitForManifestUpdated(prev_updates + 1));
+    ASSERT_NO_FATAL_FAILURE(ExpectManifestEquals({}, runner.ManifestId()));
+    CheckManifestUpdateRecorded(std::vector<metrics::ManifestUpdateData>{
+        GetManifestUpdateData(metrics::UpdateTrigger::kRunningUpdateAll,
+                              absl::StatusCode::kNotFound, 0, 0, 0, 0, 0, 0)});
+  }
+
+  {
+    SCOPED_TRACE("Create a file in place of the directory");
+
+    uint32_t prev_updates = num_manifest_updates_;
+    EXPECT_OK(path::WriteFile(cfg_.src_dir, kData, kDataSize));
+    ASSERT_TRUE(WaitForManifestUpdated(prev_updates + 2));
+    ASSERT_NO_FATAL_FAILURE(ExpectManifestEquals({}, runner.ManifestId()));
+    metrics::ManifestUpdateData update_data = GetManifestUpdateData(
+        metrics::UpdateTrigger::kRunningUpdateAll,
+        absl::StatusCode::kFailedPrecondition, 0, 0, 0, 0, 0, 0);
+    CheckManifestUpdateRecorded(
+        std::vector<metrics::ManifestUpdateData>{update_data, update_data});
+    CheckMultiSessionStartNotRecorded();
+  }
 
   EXPECT_OK(runner.Status());
   EXPECT_OK(runner.Shutdown());