Add snapshotContainer API

snapshotContainer creates a "full container" snapshot, which includes
the entire disk/filesystem and (optionally) the state of the container
memory when the snapshot is created.

Unlike directory snapshots, only a single container snapshot can be
restored (but multiple directory snapshots can be restored alongside a
container snapshot).

Author: gpanders

src/workerd/api/container.c++

@@ -10,8 +10,6 @@
 
 #include <kj/filesystem.h>
 
-#include <cmath>
-
 namespace workerd::api {
 
 namespace {
@@ -99,14 +97,16 @@ void Container::start(jsg::Lock& js, jsg::Optional<StartupOptions> maybeOptions)
   }
 
   if (!flags.getWorkerdExperimental()) {
-    JSG_REQUIRE(options.snapshots == kj::none, Error,
+    JSG_REQUIRE(options.directorySnapshots == kj::none, Error,
         "Container snapshot restore requires the 'experimental' compatibility flag.");
+    JSG_REQUIRE(options.containerSnapshot == kj::none, Error,
+        "Container full snapshot restore requires the 'experimental' compatibility flag.");
   } else {
-    KJ_IF_SOME(snapshots, options.snapshots) {
-      auto list = req.initSnapshots(snapshots.size());
-      for (auto i: kj::indices(snapshots)) {
+    KJ_IF_SOME(directorySnapshots, options.directorySnapshots) {
+      auto list = req.initDirectorySnapshots(directorySnapshots.size());
+      for (auto i: kj::indices(directorySnapshots)) {
         auto entry = list[i];
-        auto& restore = snapshots[i];
+        auto& restore = directorySnapshots[i];
         auto& snap = restore.snapshot;
         auto effectiveRestoreDir = snap.dir.asPtr();
         KJ_IF_SOME(mp, restore.mountPoint) {
@@ -116,13 +116,8 @@ void Container::start(jsg::Lock& js, jsg::Optional<StartupOptions> maybeOptions)
         JSG_REQUIRE_NONNULL(parseRestorePath(effectiveRestoreDir), Error,
             "Directory snapshot cannot be restored to root directory.");
 
-        double size = snap.size;
-        JSG_REQUIRE(std::isfinite(size) && size >= 0 &&
-                size <= static_cast<double>(jsg::MAX_SAFE_INTEGER) && std::floor(size) == size,
-            RangeError, "Snapshot size must be a non-negative integer <= Number.MAX_SAFE_INTEGER");
         auto snapshotBuilder = entry.initSnapshot();
         snapshotBuilder.setId(snap.id);
-        snapshotBuilder.setSize(static_cast<uint64_t>(size));
         snapshotBuilder.setDir(snap.dir);
         KJ_IF_SOME(name, snap.name) {
           snapshotBuilder.setName(name);
@@ -132,6 +127,14 @@ void Container::start(jsg::Lock& js, jsg::Optional<StartupOptions> maybeOptions)
         }
       }
     }
+
+    KJ_IF_SOME(containerSnapshot, options.containerSnapshot) {
+      auto builder = req.initContainerSnapshot();
+      builder.setId(containerSnapshot.id);
+      KJ_IF_SOME(name, containerSnapshot.name) {
+        builder.setName(name);
+      }
+    }
   }
 
   req.setCompatibilityFlags(flags);
@@ -160,17 +163,45 @@ jsg::Promise<Container::DirectorySnapshot> Container::snapshotDirectory(
       .then(
           js, [](jsg::Lock& js, capnp::Response<rpc::Container::SnapshotDirectoryResults> results) {
     auto snapshot = results.getSnapshot();
+    JSG_REQUIRE(snapshot.getSize() <= jsg::MAX_SAFE_INTEGER, RangeError,
+        "Snapshot size exceeds Number.MAX_SAFE_INTEGER");
+
     jsg::Optional<kj::String> name = kj::none;
-    auto snapshotName = snapshot.getName();
-    if (snapshotName.size() > 0) {
-      name = kj::str(snapshotName);
+    if (snapshot.getName().size() > 0) {
+      name = kj::str(snapshot.getName());
     }
 
+    return Container::DirectorySnapshot{kj::str(snapshot.getId()),
+      static_cast<double>(snapshot.getSize()), kj::str(snapshot.getDir()), kj::mv(name)};
+  });
+}
+
+jsg::Promise<Container::Snapshot> Container::snapshotContainer(
+    jsg::Lock& js, SnapshotOptions options) {
+  JSG_REQUIRE(
+      running, Error, "snapshotContainer() cannot be called on a container that is not running.");
+
+  auto req = rpcClient->snapshotContainerRequest();
+
+  KJ_IF_SOME(name, options.name) {
+    req.setName(name);
+  }
+
+  return IoContext::current()
+      .awaitIo(js, req.send())
+      .then(
+          js, [](jsg::Lock& js, capnp::Response<rpc::Container::SnapshotContainerResults> results) {
+    auto snapshot = results.getSnapshot();
     JSG_REQUIRE(snapshot.getSize() <= jsg::MAX_SAFE_INTEGER, RangeError,
         "Snapshot size exceeds Number.MAX_SAFE_INTEGER");
 
-    return Container::DirectorySnapshot{kj::str(snapshot.getId()),
-      static_cast<double>(snapshot.getSize()), kj::str(snapshot.getDir()), kj::mv(name)};
+    jsg::Optional<kj::String> name = kj::none;
+    if (snapshot.getName().size() > 0) {
+      name = kj::str(snapshot.getName());
+    }
+
+    return Container::Snapshot{
+      kj::str(snapshot.getId()), static_cast<double>(snapshot.getSize()), kj::mv(name)};
   });
 }
 

src/workerd/api/container.h

@@ -41,16 +41,49 @@ class Container: public jsg::Object {
     jsg::Optional<kj::String> name;
 
     JSG_STRUCT(dir, name);
+    JSG_STRUCT_TS_OVERRIDE_DYNAMIC(CompatibilityFlags::Reader flags) {
+      if (!flags.getWorkerdExperimental()) {
+        JSG_TS_OVERRIDE(type DirectorySnapshotOptions = never);
+      }
+    }
   };
 
-  struct SnapshotRestoreParams {
+  struct DirectorySnapshotRestoreParams {
     DirectorySnapshot snapshot;
     jsg::Optional<kj::String> mountPoint;
 
     JSG_STRUCT(snapshot, mountPoint);
     JSG_STRUCT_TS_OVERRIDE_DYNAMIC(CompatibilityFlags::Reader flags) {
       if (!flags.getWorkerdExperimental()) {
-        JSG_TS_OVERRIDE(type SnapshotRestoreParams = never);
+        JSG_TS_OVERRIDE(type DirectorySnapshotRestoreParams = never);
+      }
+    }
+  };
+
+  struct Snapshot {
+    kj::String id;
+    double size;
+    jsg::Optional<kj::String> name;
+
+    JSG_STRUCT(id, size, name);
+    JSG_STRUCT_TS_OVERRIDE_DYNAMIC(CompatibilityFlags::Reader flags) {
+      if (!flags.getWorkerdExperimental()) {
+        JSG_TS_OVERRIDE(type Snapshot = never);
+      }
+    }
+  };
+
+  struct SnapshotOptions {
+    jsg::Optional<kj::String> name;
+
+    JSG_STRUCT(name);
+    JSG_STRUCT_TS_OVERRIDE_DYNAMIC(CompatibilityFlags::Reader flags) {
+      if (flags.getWorkerdExperimental()) {
+        JSG_TS_OVERRIDE(ContainerSnapshotOptions {
+          name?: string;
+        });
+      } else {
+        JSG_TS_OVERRIDE(type SnapshotOptions = never);
       }
     }
   };
@@ -61,11 +94,18 @@ class Container: public jsg::Object {
     jsg::Optional<jsg::Dict<kj::String>> env;
     jsg::Optional<int64_t> hardTimeout;
     jsg::Optional<jsg::Dict<kj::String>> labels;
-    jsg::Optional<kj::Array<SnapshotRestoreParams>> snapshots;
+    jsg::Optional<kj::Array<DirectorySnapshotRestoreParams>> directorySnapshots;
+    jsg::Optional<Snapshot> containerSnapshot;
 
     // TODO(containers): Allow intercepting stdin/stdout/stderr by specifying streams here.
 
-    JSG_STRUCT(entrypoint, enableInternet, env, hardTimeout, labels, snapshots);
+    JSG_STRUCT(entrypoint,
+        enableInternet,
+        env,
+        hardTimeout,
+        labels,
+        directorySnapshots,
+        containerSnapshot);
     JSG_STRUCT_TS_OVERRIDE_DYNAMIC(CompatibilityFlags::Reader flags) {
       if (flags.getWorkerdExperimental()) {
         JSG_TS_OVERRIDE(ContainerStartupOptions {
@@ -74,7 +114,8 @@ class Container: public jsg::Object {
           env?: Record<string, string>;
           hardTimeout?: number | bigint;
           labels?: Record<string, string>;
-          snapshots?: ContainerSnapshotRestoreParams[];
+          directorySnapshots?: ContainerDirectorySnapshotRestoreParams[];
+          containerSnapshot?: ContainerSnapshot;
         });
       } else {
         JSG_TS_OVERRIDE(ContainerStartupOptions {
@@ -83,7 +124,8 @@ class Container: public jsg::Object {
           env?: Record<string, string>;
           hardTimeout?: never;
           labels?: Record<string, string>;
-          snapshots?: never;
+          directorySnapshots?: never;
+          containerSnapshot?: never;
         });
       }
     }
@@ -107,6 +149,7 @@ class Container: public jsg::Object {
       jsg::Lock& js, kj::String addr, jsg::Ref<Fetcher> binding);
   jsg::Promise<DirectorySnapshot> snapshotDirectory(
       jsg::Lock& js, DirectorySnapshotOptions options);
+  jsg::Promise<Snapshot> snapshotContainer(jsg::Lock& js, SnapshotOptions options);
 
   // TODO(containers): listenTcp()
 
@@ -124,6 +167,7 @@ class Container: public jsg::Object {
     if (flags.getWorkerdExperimental()) {
       JSG_METHOD(interceptOutboundHttps);
       JSG_METHOD(snapshotDirectory);
+      JSG_METHOD(snapshotContainer);
     }
   }
 
@@ -147,6 +191,7 @@ class Container: public jsg::Object {
 
 #define EW_CONTAINER_ISOLATE_TYPES                                                                 \
   api::Container, api::Container::DirectorySnapshot, api::Container::DirectorySnapshotOptions,     \
-      api::Container::SnapshotRestoreParams, api::Container::StartupOptions
+      api::Container::DirectorySnapshotRestoreParams, api::Container::Snapshot,                    \
+      api::Container::SnapshotOptions, api::Container::StartupOptions
 
 }  // namespace workerd::api

src/workerd/io/container.capnp

@@ -48,16 +48,19 @@ interface Container @0x9aaceefc06523bca {
     labels @5 :List(Label);
     # Optional key-value metadata labels for metrics/observability.
 
-    snapshots @6 :List(SnapshotRestoreParams);
+    directorySnapshots @6 :List(DirectorySnapshotRestoreParams);
     # Directory snapshots to restore before the container starts.
+
+    containerSnapshot @7 :ContainerSnapshot;
+    # Full container snapshot to restore before the container starts.
   }
 
   struct Label {
     name @0 :Text;
     value @1 :Text;
   }
 
-  struct SnapshotRestoreParams {
+  struct DirectorySnapshotRestoreParams {
     snapshot @0 :DirectorySnapshot;
     # The snapshot to restore.
 
@@ -90,6 +93,24 @@ interface Container @0x9aaceefc06523bca {
     # Optional human-friendly name. Empty string means not set.
   }
 
+  struct ContainerSnapshot {
+    # Opaque handle to a full container snapshot.
+
+    id @0 :Text;
+    # Unique identifier of the snapshot.
+
+    size @1 :UInt64;
+    # Snapshot size, in bytes.
+
+    name @2 :Text;
+    # Optional human-friendly name. Empty string means not set.
+  }
+
+  struct SnapshotContainerParams {
+    name @0 :Text;
+    # Optional human-friendly name. Empty string means not set.
+  }
+
   monitor @2 () -> (exitCode: Int32);
   # Waits for the container to shut down.
   #
@@ -181,4 +202,7 @@ interface Container @0x9aaceefc06523bca {
 
   snapshotDirectory @10 SnapshotDirectoryParams -> (snapshot :DirectorySnapshot);
   # Creates a snapshot for a directory in the running container.
+
+  snapshotContainer @11 SnapshotContainerParams -> (snapshot :ContainerSnapshot);
+  # Creates a full container snapshot for the running container.
 }

src/workerd/server/container-client.c++

@@ -1524,15 +1524,18 @@ kj::Promise<void> ContainerClient::start(StartContext context) {
 
   internetEnabled = params.getEnableInternet();
 
+  JSG_REQUIRE(!params.hasContainerSnapshot(), Error,
+      "container.start({ containerSnapshot }) is not implemented for local Docker containers yet.");
+
   // If startup fails after we clone any snapshot volumes, tear down the app container first and
   // then delete those clone volumes so we don't leave mounted Docker volumes behind.
   KJ_DEFER(if (!containerStarted.load(std::memory_order_acquire)) {
     waitUntilTasks.add(destroyContainer().attach(addRef()));
   });
 
   kj::Vector<SnapshotRestoreMount> restoreMounts;
-  if (params.hasSnapshots()) {
-    auto snapshotList = params.getSnapshots();
+  if (params.hasDirectorySnapshots()) {
+    auto snapshotList = params.getDirectorySnapshots();
     restoreMounts.reserve(snapshotList.size());
     for (auto i: kj::zeroTo(snapshotList.size())) {
       auto entry = snapshotList[i];
@@ -1724,6 +1727,17 @@ kj::Promise<void> ContainerClient::snapshotDirectory(SnapshotDirectoryContext co
   }
 }
 
+kj::Promise<void> ContainerClient::snapshotContainer(SnapshotContainerContext context) {
+  auto [ready, done] = getRpcTurn();
+  co_await ready;
+  KJ_DEFER(done->fulfill());
+
+  JSG_REQUIRE(containerStarted.load(std::memory_order_acquire), Error,
+      "snapshotContainer() requires a running container.");
+  JSG_FAIL_REQUIRE(
+      Error, "snapshotContainer() is not implemented for local Docker containers yet.");
+}
+
 kj::Promise<void> ContainerClient::getTcpPort(GetTcpPortContext context) {
   co_await mutationQueue.addBranch();
 

src/workerd/server/container-client.h

@@ -77,6 +77,7 @@ class ContainerClient final: public rpc::Container::Server, public kj::Refcounte
   kj::Promise<void> setEgressHttp(SetEgressHttpContext context) override;
   kj::Promise<void> setEgressHttps(SetEgressHttpsContext context) override;
   kj::Promise<void> snapshotDirectory(SnapshotDirectoryContext context) override;
+  kj::Promise<void> snapshotContainer(SnapshotContainerContext context) override;
 
   kj::Own<ContainerClient> addRef();