cluster: call Durable Objects' drain() before a draining instance lets them go

src/workerd/server/cluster/cluster.rs

@@ -180,8 +180,13 @@ impl Cluster {
         }
         Ok(match claim {
             Claim::Owned(owned) => {
-                let epoch = owned.epoch;
                 let key = format!("{ns}/{actor}");
+                // The node may have started draining while the claim awaited the bucket: then
+                // its objects are being let go, and one opened now would outlive that.
+                if !self.ownership.may_keep.get() {
+                    return Err(crate::ownership::draining(&key));
+                }
+                let epoch = owned.epoch;
                 self.claims.borrow_mut().insert(key, owned);
                 Route::Owned(epoch)
             }
@@ -388,11 +393,20 @@ impl Cluster {
         }
     }
 
+    /// Stops claiming before the claim reaches the bucket, but keeps the
+    /// lease: the objects this node owns stay its own, their commits still
+    /// acknowledged, while they flush before `start_draining`. One its record
+    /// names that is not open here still opens here.
+    pub fn stop_claiming(&self) {
+        self.ownership.may_claim.set(false);
+    }
+
     /// Stops claiming before the claim reaches the bucket, and parks the
     /// lease loop; then `deregister`.
     pub fn start_draining(&self) {
         self.lease.borrow_mut().draining = true;
         self.ownership.may_claim.set(false);
+        self.ownership.may_keep.set(false);
     }
 
     /// The first renewal, before the node serves.
@@ -450,12 +464,138 @@ async fn check_conditional_writes(bucket: &Bucket, node: &str) -> Result<()> {
 
 #[cfg(test)]
 mod tests {
+    use std::cell::Cell;
+    use std::cell::RefCell;
+    use std::rc::Rc;
+
+    use cluster_logic::lease;
     use futures::executor::block_on;
 
+    use super::Cluster;
+    use super::Route;
+    use super::Setup;
     use super::check_conditional_writes;
+    use super::now_ms;
     use crate::Bucket;
     use crate::bucket::memory::Memory;
 
+    /// A memory bucket that runs `hook` as it serves the `nth` read of an
+    /// owner record, from 1.
+    struct Hooked {
+        inner: Memory,
+        nth: Cell<usize>,
+        hook: RefCell<Option<Box<dyn FnOnce()>>>,
+    }
+
+    #[async_trait::async_trait(?Send)]
+    impl crate::Store for Rc<Hooked> {
+        async fn get(
+            &self,
+            key: &str,
+            to: Option<&std::path::Path>,
+        ) -> crate::Result<Option<(Vec<u8>, crate::Version)>> {
+            if key.ends_with("own.json") {
+                self.nth.set(self.nth.get().saturating_sub(1));
+                if self.nth.get() == 0
+                    && let Some(hook) = self.hook.borrow_mut().take()
+                {
+                    hook();
+                }
+            }
+            self.inner.get(key, to).await
+        }
+
+        async fn put(
+            &self,
+            key: &str,
+            body: crate::bucket::Payload,
+            guard: Option<crate::Guard>,
+            metadata: &[(&str, String)],
+        ) -> crate::Result<Option<crate::Version>> {
+            self.inner.put(key, body, guard, metadata).await
+        }
+
+        async fn list(&self, prefix: &str) -> crate::Result<Vec<crate::bucket::Entry>> {
+            self.inner.list(prefix).await
+        }
+
+        async fn delete(&self, key: &str) -> crate::Result<()> {
+            self.inner.delete(key).await
+        }
+    }
+
+    fn node(bucket: &Bucket, name: &str) -> Rc<Cluster> {
+        let setup = Setup {
+            url: String::new(),
+            options: Vec::new(),
+            node: name.to_owned(),
+            address: format!("{name}:1"),
+            ttl_ms: lease::LEASE_MS,
+            key: Vec::new(),
+        };
+        Cluster::over(bucket.clone(), setup)
+    }
+
+    #[test]
+    fn a_node_that_stopped_claiming_keeps_its_objects_and_its_lease() {
+        block_on(async {
+            let bucket = Bucket::memory("");
+            let (a, b) = (node(&bucket, "a"), node(&bucket, "b"));
+            a.join().await.unwrap();
+            b.join().await.unwrap();
+            let kept = a.acquire_or_route("ns", "kept", now_ms()).await.unwrap();
+            assert!(matches!(kept, Route::Owned(_)));
+
+            a.stop_claiming();
+            // Nothing new, nor anything it would have to take over...
+            assert!(a.acquire_or_route("ns", "new", now_ms()).await.is_err());
+            // ...but the lease loop is not parked, and the lease still renews,
+            assert!(!a.lease.borrow().draining);
+            a.renew(now_ms()).await.unwrap();
+            assert!(a.lease_holds());
+            // so a peer routes to what it owns rather than claiming it,
+            match b.acquire_or_route("ns", "kept", now_ms()).await.unwrap() {
+                Route::Remote { node, .. } => assert_eq!(node, "a"),
+                Route::Owned(_) => panic!("a peer claimed an object of a live node"),
+            }
+            // and it still takes up what it owns.
+            let again = a.acquire_or_route("ns", "kept", now_ms()).await.unwrap();
+            assert!(matches!(again, Route::Owned(_)));
+
+            // Draining, it lets its objects go, its own too.
+            a.start_draining();
+            assert!(a.lease.borrow().draining);
+            assert!(a.acquire_or_route("ns", "kept", now_ms()).await.is_err());
+        });
+    }
+
+    #[test]
+    fn a_node_that_starts_draining_mid_claim_does_not_keep_the_object() {
+        block_on(async {
+            let hooked = Rc::new(Hooked {
+                inner: Memory::default(),
+                nth: Cell::new(usize::MAX),
+                hook: RefCell::new(None),
+            });
+            let bucket = Bucket::new(Rc::clone(&hooked), "");
+            let a = node(&bucket, "a");
+            a.join().await.unwrap();
+            let kept = a.acquire_or_route("ns", "kept", now_ms()).await.unwrap();
+            assert!(matches!(kept, Route::Owned(_)));
+            // Not open here any more: taking it up again moves it to the
+            // next epoch.
+            a.close_actor("ns", "kept");
+            a.stop_claiming();
+
+            // That reads its record twice: to find it is this node's, then
+            // to advance it. The drain starts between the two.
+            let drainer = Rc::clone(&a);
+            *hooked.hook.borrow_mut() = Some(Box::new(move || drainer.start_draining()));
+            hooked.nth.set(2);
+            assert!(a.acquire_or_route("ns", "kept", now_ms()).await.is_err());
+        });
+    }
+
     #[test]
     fn a_bucket_that_ignores_preconditions_is_refused() {
         block_on(async {

src/workerd/server/cluster/ownership.rs

@@ -79,6 +79,9 @@ pub struct Ownership {
     registry: Rc<Registry>,
     /// Whether this instance may claim objects: off once it drains.
     pub may_claim: Cell<bool>,
+    /// Whether this instance still takes up the objects its owner records
+    /// name, while it holds its lease: off once it lets them go.
+    pub may_keep: Cell<bool>,
     /// Where each routed object's owner was, and when that was read.
     routes: RefCell<HashMap<String, Route>>,
     /// The record each claim wrote, and the guard that it is still the
@@ -105,7 +108,7 @@ const ROUTE_REFRESH_MS: u64 = 5_000;
 /// How many times a claim retries when another instance keeps winning.
 const ATTEMPTS: usize = 8;
 
-fn draining(key: &str) -> Error {
+pub fn draining(key: &str) -> Error {
     Error::Replica(format!("{key}: this node is draining"))
 }
 
@@ -119,6 +122,7 @@ impl Ownership {
             bucket,
             registry,
             may_claim: Cell::new(true),
+            may_keep: Cell::new(true),
             routes: RefCell::new(HashMap::new()),
             claimed: RefCell::new(HashMap::new()),
         }
@@ -163,8 +167,9 @@ impl Ownership {
                 Some((body, version)) => {
                     let record: OwnerRecord = serde_json::from_slice(&body)?;
                     if record.node == me {
-                        // A draining node lets its objects go, its own too.
-                        if !self.may_claim.get() {
+                        // A node that stopped claiming still wakes what it
+                        // owns, until it lets its objects go.
+                        if !self.may_keep.get() {
                             return Err(draining(&key));
                         }
                         return Ok(Claim::Owned(record.owned()));

src/workerd/server/factory/worker-factory-actor.c++

@@ -5,6 +5,7 @@
 #include "worker-factory-impl.h"
 
 #include <workerd/api/actor-state.h>
+#include <workerd/api/global-scope.h>
 #include <workerd/io/actor-cache.h>
 #include <workerd/io/io-context.h>
 #include <workerd/server/actor-id-impl.h>
@@ -458,6 +459,80 @@ class ActorHooksImpl final: public Worker::Actor::Loopback,
   }
 };
 
+// A draining cluster instance's call of an actor's `drain()` method, the actor's chance to flush
+// what it keeps in memory while the instance still owns it. A class without the method gets nothing
+// (`ran` stays false). The event completes once what the method wrote is durable: in a cluster, a
+// commit is durable once the bucket has it.
+class DrainCustomEvent final: public WorkerInterface::CustomEvent {
+ public:
+  explicit DrainCustomEvent(bool& ran): ran(ran) {}
+
+  kj::Promise<Result> run(kj::Own<IoContext_IncomingRequest> incomingRequest,
+      kj::Maybe<kj::StringPtr> entrypointName,
+      kj::Maybe<Worker::VersionInfo> versionInfo,
+      Frankenvalue props,
+      kj::TaskSet& waitUntilTasks,
+      bool isDynamicDispatch) override {
+    auto& context = incomingRequest->getContext();
+    incomingRequest->delivered("drain"_kj);
+    KJ_DEFER({ incomingRequest->drain(waitUntilTasks, kj::mv(incomingRequest)); });
+
+    EventOutcome outcome = EventOutcome::OK;
+    try {
+      co_await context.run([this, entrypointName, versionInfo = kj::mv(versionInfo),
+                               props = kj::mv(props), isDynamicDispatch](
+                               Worker::Lock& lock, IoContext& ioctx) mutable -> kj::Promise<void> {
+        KJ_IF_SOME(handler,
+            lock.getExportedHandler(entrypointName, kj::mv(versionInfo), kj::mv(props),
+                ioctx.getActor(), isDynamicDispatch)) {
+          jsg::Lock& js = lock;
+          auto self = jsg::JsObject(handler->self.getHandle(js));
+          KJ_IF_SOME(method, self.get(js, "drain"_kj).tryCast<jsg::JsFunction>()) {
+            ran = true;
+            v8::LocalVector<v8::Value> args(js.v8Isolate);
+            return ioctx.awaitJs(js, js.tryCatch([&] {
+              return js.toVoidPromise(method.call(js, self, args));
+            }, [&](jsg::Value exception) { return js.rejectedPromise<void>(kj::mv(exception)); }));
+          }
+        }
+        return kj::READY_NOW;
+      });
+      // Nothing to wait for from an actor whose method did not run.
+      if (ran) co_await context.waitForOutputLocks();
+    } catch (kj::Exception& e) {
+      incomingRequest->getMetrics().reportFailure(e);
+      context.logUncaughtExceptionAsync(
+          UncaughtExceptionSource::ASYNC_TASK, e.clone(), *incomingRequest);
+      outcome = EventOutcome::EXCEPTION;
+    }
+    co_return Result{.outcome = outcome};
+  }
+
+  kj::Promise<Result> sendRpc(capnp::HttpOverCapnpFactory&,
+      capnp::ByteStreamFactory&,
+      FrankenvalueHandler&,
+      rpc::EventDispatcher::Client) override {
+    // An instance drains only the actors it runs.
+    KJ_FAIL_REQUIRE("a drain does not cross between instances");
+  }
+
+  kj::Promise<Result> notSupported() override {
+    KJ_FAIL_REQUIRE("the target does not take a drain");
+  }
+
+  uint16_t getType() override {
+    // Defined here: the event never leaves workerd.
+    return 0x8001;
+  }
+
+  tracing::EventInfo getEventInfo() const override {
+    return tracing::CustomEventInfo();
+  }
+
+ private:
+  bool& ran;
+};
+
 // Schedules a root actor's alarms with the namespace's scheduler, at once, and in a cluster indexes
 // them after `priorTask`, as production's alarm manager orders its requests: the index takes them
 // in the order they were made.
@@ -773,6 +848,59 @@ kj::Maybe<kj::Rc<HibernationManager>> actor_hibernation_manager(const ActorHandl
       [](HibernationManager& m) { return kj::Rc<HibernationManager>(m.addRef()); });
 }
 
+kj::Promise<DrainMethod> actor_drain_method(const ActorHandle& handle) {
+  auto& actor = handle.getActor();
+  auto asyncLock = co_await actor.getWorker().takeAsyncLockWithoutRequest(nullptr);
+  // An actor still being constructed, or whose construction failed, can only tell in a request.
+  auto method = DrainMethod::Unknown;
+  actor.getWorker().runInLockScope(asyncLock, [&](Worker::Lock& lock) {
+    JSG_WITHIN_CONTEXT_SCOPE(lock, lock.getContext(), [&](jsg::Lock& js) {
+      kj::runCatchingExceptions([&]() {
+        // The constructed instance; outside a request, anything else throws.
+        KJ_IF_SOME(handler,
+            lock.getExportedHandler(kj::none, kj::none, Frankenvalue(), actor, false)) {
+          // Only a data property that is a function counts, found along the prototype chain
+          // without running the app's code: an accessor or a Proxy leaves it to the request.
+          method = js.tryCatch([&] {
+            auto context = js.v8Context();
+            v8::Local<v8::String> name = js.strIntern("drain"_kj);
+            v8::Local<v8::String> get = js.strIntern("get"_kj);
+            v8::Local<v8::String> set = js.strIntern("set"_kj);
+            v8::Local<v8::String> value = js.strIntern("value"_kj);
+            v8::Local<v8::Value> current = handler->self.getHandle(js);
+            for (auto depth = 0; depth < 64 && current->IsObject(); ++depth) {
+              auto object = current.As<v8::Object>();
+              if (object->IsProxy()) return DrainMethod::Unknown;
+              auto descriptor = jsg::check(object->GetOwnPropertyDescriptor(context, name));
+              if (descriptor->IsObject()) {
+                auto fields = descriptor.As<v8::Object>();
+                if (jsg::check(fields->HasOwnProperty(context, get)) ||
+                    jsg::check(fields->HasOwnProperty(context, set))) {
+                  return DrainMethod::Unknown;
+                }
+                return jsg::check(fields->Get(context, value))->IsFunction() ? DrainMethod::Some
+                                                                             : DrainMethod::None;
+              }
+              current = object->GetPrototype();
+            }
+            return current->IsNull() ? DrainMethod::None : DrainMethod::Unknown;
+          }, [&](jsg::Value) { return DrainMethod::Unknown; });
+        } else {
+          method = DrainMethod::None;
+        }
+      });
+    });
+  });
+  co_return method;
+}
+
+kj::Promise<bool> worker_drain(kj::Own<WorkerInterface> worker) {
+  bool ran = false;
+  auto result = co_await worker->customEvent(kj::heap<DrainCustomEvent>(ran));
+  KJ_REQUIRE(result.outcome == EventOutcome::OK, "the Durable Object's drain() failed");
+  co_return ran;
+}
+
 kj::Promise<FacetStartInfo> facet_start_resolve(kj::Own<FacetStart> start) {
   auto info = co_await (*start)();
   co_await info.ensureAllResolved();

src/workerd/server/factory/worker-factory.h

@@ -99,6 +99,7 @@ struct ActorStorageSpec;
 struct HeaderEdit;
 struct UdpDatagram;
 enum class AlarmElsewhere : uint8_t;
+enum class DrainMethod : uint8_t;
 
 // How the config message is read: configs can legitimately be very large and are not malicious.
 constexpr capnp::ReaderOptions CONFIG_READER_OPTIONS = {.traversalLimitInWords = kj::maxValue};
@@ -217,6 +218,8 @@ kj::Own<WorkerInterface> worker_interface_attach(
     kj::Own<WorkerInterface> inner, ::rust::Box<KeepAlive> keep);
 kj::Promise<::rust::Vec<uint64_t>> worker_bench(
     kj::Own<WorkerInterface> worker, ::rust::Slice<const uint64_t> params);
+kj::Promise<bool> worker_drain(kj::Own<WorkerInterface> worker);
+kj::Promise<DrainMethod> actor_drain_method(const ActorHandle& actor);
 kj::Own<ActorHandle> worker_new_actor(const CompiledWorker& worker,
     ::rust::Str className,
     kj::Maybe<kj::Own<Frankenvalue>> props,