feat(server): stop when the host asks

.github/workflows/ci.yml

@@ -41,6 +41,8 @@ jobs:
         run: |
           cmake -B build \
             -DCMAKE_TOOLCHAIN_FILE="$VCPKG_INSTALLATION_ROOT/scripts/buildsystems/vcpkg.cmake" \
+            -DVCPKG_TARGET_TRIPLET=x64-linux-release \
+            -DVCPKG_HOST_TRIPLET=x64-linux-release \
             -DCMAKE_BUILD_TYPE=Release \
             -DVCPKG_MANIFEST_FEATURES=tests \
             -DGO_PLUGIN_LOG_ABSL=ON \

CMakePresets.json

@@ -6,7 +6,9 @@
       "generator": "Ninja",
       "binaryDir": "${sourceDir}/build",
       "cacheVariables": {
-        "CMAKE_TOOLCHAIN_FILE": "$penv{VCPKG_ROOT}/scripts/buildsystems/vcpkg.cmake"
+        "CMAKE_TOOLCHAIN_FILE": "$penv{VCPKG_ROOT}/scripts/buildsystems/vcpkg.cmake",
+        "VCPKG_TARGET_TRIPLET": "x64-linux-release",
+        "VCPKG_HOST_TRIPLET": "x64-linux-release"
       }
     }
   ]

README.md

@@ -20,6 +20,9 @@ This library handles:
   configured automatically.
 - **Handshake line output** – writes the correctly formatted line to stdout so the host can connect.
 - **Health-check service** – the built-in gRPC health-check service is registered automatically.
+- **Stopping on request** – go-plugin's controller is served, so a host asking the plugin to stop shuts the server
+  down and returns `Wait()`, letting `main` unwind instead of being killed after the host's grace period.
+- **Parent-death watch** – if the host process dies without asking, the server shuts down on its own.
 - **Logging the host can read** – see below.
 
 ## Logging

example/host/generate.go

@@ -1,3 +1,3 @@
 package main
 
-//go:generate ../../build/vcpkg_installed/x64-linux/tools/protobuf/protoc --proto_path=../proto --go_out=proto --go_opt=paths=source_relative --go-grpc_out=proto --go-grpc_opt=paths=source_relative ../proto/greeter.proto
+//go:generate ../../build/vcpkg_installed/x64-linux-release/tools/protobuf/protoc --proto_path=../proto --go_out=proto --go_opt=paths=source_relative --go-grpc_out=proto --go-grpc_opt=paths=source_relative ../proto/greeter.proto

include/go_plugin/server.hpp

@@ -79,8 +79,10 @@ class PluginServer {
 
     /**
      * Validates the magic cookie, picks a port, starts the gRPC server,
-     * registers services and the built-in health-check service, then writes
-     * the go-plugin handshake line to the configured output stream.
+     * registers services, the built-in health-check service and go-plugin's
+     * controller, then writes the go-plugin handshake line to the configured
+     * output stream. The controller answers the host's request to stop by
+     * shutting the server down, which returns Wait().
      *
      * @param error  If non-null, receives a human-readable error message on
      *               failure.
@@ -89,8 +91,9 @@ class PluginServer {
     bool Start(std::string* error = nullptr);
 
     /**
-     * Blocks until the server stops (i.e. until Shutdown() is called or the
-     * host process exits — see the parent-death watchdog started by Start()).
+     * Blocks until the server stops: Shutdown() is called, the host asks the
+     * plugin to stop, or the host process exits (see the parent-death watchdog
+     * started by Start()).
      */
     void Wait();
 
@@ -101,23 +104,30 @@ class PluginServer {
     int port() const { return port_; }
 
 private:
-    // Watches for the host (parent) process disappearing. go-plugin hosts stop a
-    // plugin by killing it, but when the host itself dies unexpectedly (e.g. it
-    // is SIGKILLed) nothing signals the plugin and Wait() would block forever,
-    // orphaning the process. The watchdog detects reparenting (getppid changes)
-    // and calls Shutdown(), mirroring the Go go-plugin server's built-in
-    // orphan protection. Started by Start(), stopped by Shutdown()/destruction.
+    // Watches for the host (parent) process disappearing. When the host dies
+    // unexpectedly (e.g. it is SIGKILLed) nothing signals the plugin and Wait()
+    // would block forever, orphaning the process. The watchdog detects
+    // reparenting (getppid changes) and calls Shutdown(), mirroring the Go
+    // go-plugin server's built-in orphan protection. It also carries out the
+    // host's request to stop: the request arrives on a gRPC thread, which
+    // cannot wait for the server it runs in to shut down. Started by Start(),
+    // stopped by Shutdown()/destruction.
     void StartParentWatchdog();
     void StopParentWatchdog();
 
+    // Called by the controller when the host asks the plugin to stop.
+    void RequestShutdown();
+
     ServeConfig config_;
+    std::unique_ptr<grpc::Service> controller_;
     std::unique_ptr<grpc::Server> server_;
     int port_ = 0;
 
     std::thread parent_watchdog_;
     std::mutex watchdog_mu_;
     std::condition_variable watchdog_cv_;
     bool watchdog_stop_ = false;
+    bool shutdown_requested_ = false;
     pid_t host_pid_ = 0;
 };
 

src/server.cpp

@@ -6,6 +6,7 @@
 #include <chrono>
 #include <cstdlib>
 #include <cstring>
+#include <functional>
 #include <iostream>
 #include <numeric>
 #include <random>
@@ -21,11 +22,36 @@
 
 #include <grpcpp/grpcpp.h>
 #include <grpcpp/health_check_service_interface.h>
+#include <grpcpp/impl/rpc_service_method.h>
+#include <grpcpp/support/method_handler.h>
 
 namespace go_plugin {
 
 namespace {
 
+// go-plugin's GRPCController: the service a host calls to ask the plugin to
+// stop. Its one method takes and returns plugin.Empty, which is zero bytes on
+// the wire, so it is served over raw byte buffers and needs no generated code.
+class ControllerService final : public grpc::Service {
+public:
+  explicit ControllerService(std::function<void()> on_shutdown) : on_shutdown_(std::move(on_shutdown)) {
+    using Handler = grpc::internal::RpcMethodHandler<ControllerService, grpc::ByteBuffer, grpc::ByteBuffer>;
+    AddMethod(new grpc::internal::RpcServiceMethod(
+        "/plugin.GRPCController/Shutdown", grpc::internal::RpcMethod::NORMAL_RPC,
+        new Handler(
+            [](ControllerService *service, grpc::ServerContext *, const grpc::ByteBuffer *, grpc::ByteBuffer *reply) {
+              grpc::Slice empty;
+              *reply = grpc::ByteBuffer(&empty, 1);
+              service->on_shutdown_();
+              return grpc::Status::OK;
+            },
+            this)));
+  }
+
+private:
+  std::function<void()> on_shutdown_;
+};
+
 bool ValidateMagicCookie(const HandshakeConfig &h) {
   const char *val = std::getenv(h.magic_cookie_key.c_str());
   return val && std::string(val) == h.magic_cookie_value;
@@ -143,6 +169,8 @@ bool PluginServer::Start(std::string *out_error) {
   for (auto *svc : config_.services) {
     builder.RegisterService(svc);
   }
+  controller_ = std::make_unique<ControllerService>([this] { RequestShutdown(); });
+  builder.RegisterService(controller_.get());
 
   server_ = builder.BuildAndStart();
   if (!server_) {
@@ -188,6 +216,11 @@ void PluginServer::StartParentWatchdog() {
   parent_watchdog_ = std::thread([this]() {
     std::unique_lock<std::mutex> lock(watchdog_mu_);
     while (!watchdog_stop_) {
+      if (shutdown_requested_) {
+        lock.unlock();
+        Shutdown();
+        return;
+      }
       // Reparenting (parent pid changes, typically to the reaper/init) means the
       // host that launched us is gone. Drive the same Shutdown() the host would
       // have, so we don't linger as an orphan.
@@ -202,6 +235,14 @@ void PluginServer::StartParentWatchdog() {
   });
 }
 
+void PluginServer::RequestShutdown() {
+  {
+    std::lock_guard<std::mutex> lock(watchdog_mu_);
+    shutdown_requested_ = true;
+  }
+  watchdog_cv_.notify_all();
+}
+
 void PluginServer::StopParentWatchdog() {
   {
     std::lock_guard<std::mutex> lock(watchdog_mu_);

tests/test_server.cpp

@@ -4,6 +4,7 @@
 #include <thread>
 
 #include <grpcpp/grpcpp.h>
+#include <grpcpp/impl/client_unary_call.h>
 #include <gtest/gtest.h>
 
 #include "plugin_fixture.hpp"
@@ -74,5 +75,39 @@ TEST_F(Server, WaitUnblocksAfterShutdown) {
     EXPECT_TRUE(wait_returned);
 }
 
+// The host asks a plugin to stop through go-plugin's controller, and the
+// plugin's Wait() returns without anything else calling Shutdown().
+TEST_F(Server, StopsWhenTheHostAsks) {
+    std::string error;
+    ASSERT_TRUE(Start(&error)) << error;
+
+    std::atomic<bool> wait_returned{false};
+    std::thread waiter([&] {
+        server().Wait();
+        wait_returned = true;
+    });
+
+    auto channel = grpc::CreateChannel(target(), grpc::InsecureChannelCredentials());
+    grpc::ClientContext context;
+    context.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(5));
+    grpc::Slice empty;
+    grpc::ByteBuffer request(&empty, 1);
+    grpc::ByteBuffer reply;
+    grpc::Status status = grpc::internal::BlockingUnaryCall<grpc::ByteBuffer, grpc::ByteBuffer>(
+        channel.get(),
+        grpc::internal::RpcMethod("/plugin.GRPCController/Shutdown", grpc::internal::RpcMethod::NORMAL_RPC),
+        &context, request, &reply);
+    EXPECT_TRUE(status.ok()) << status.error_message();
+
+    for (int i = 0; i < 50 && !wait_returned; ++i) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(100));
+    }
+    EXPECT_TRUE(wait_returned) << "Wait must return once the host asks the plugin to stop";
+    if (!wait_returned) {
+        server().Shutdown();
+    }
+    waiter.join();
+}
+
 }  // namespace
 }  // namespace go_plugin::test