From 89eb6da13f71a7e66e643689b178c27e85eda06d Mon Sep 17 00:00:00 2001
From: Max Desiatov <m_desiatov@apple.com>
Date: Wed, 11 Dec 2024 13:14:45 +0000
Subject: [PATCH 1/2] Add `shutDown` method to `PluginMessageHandler`

This method deterministically and synchronously cleans up resources that cannot be dealt with in a deinitializer due to possible errors thrown during the clean up. Usually this includes closure of file handles, sockets, shutting down external processes and IPC resources set up for these processes, etc.

This method is required for `swift-plugin-server` to properly clean up file handles there were opened for communication with Wasm macros, as prototyped in https://github.com/swiftlang/swift/pull/73031.
---
 .../CompilerPluginMessageHandler.swift                 | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift b/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift
index 6c78e03e845..da680bb2e20 100644
--- a/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift
+++ b/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift
@@ -133,6 +133,12 @@ public class CompilerPluginMessageListener<Connection: MessageConnection, Handle
 public protocol PluginMessageHandler {
   /// Handles a single message received from the plugin host.
   func handleMessage(_ message: HostToPluginMessage) -> PluginToHostMessage
+
+  /// Deterministically and synchronously cleans up resources that cannot be dealt with in
+  /// a deinitializer due to possible errors thrown during the clean up. Usually this
+  /// includes closure of file handles, sockets, shutting down external processes and IPC
+  /// resources set up for these processes, etc.
+  func shutDown() throws
 }
 
 /// A `PluginMessageHandler` that uses a `PluginProvider`.
@@ -226,6 +232,10 @@ public class PluginProviderMessageHandler<Provider: PluginProvider>: PluginMessa
       return .loadPluginLibraryResult(loaded: diags.isEmpty, diagnostics: diags)
     }
   }
+
+  /// Empty implementation for the default message handler, since all resources are automatically
+  /// cleaned up in the synthesized initializer.
+  public func shutDown() throws {}
 }
 
 @_spi(PluginMessage)

From 76af38bd8e7ed324b92d492a4f470593613eabca Mon Sep 17 00:00:00 2001
From: Max Desiatov <m_desiatov@apple.com>
Date: Wed, 11 Dec 2024 14:56:44 +0000
Subject: [PATCH 2/2] SwiftCompilerPlugin: Call new `shutDown` function from
 `main`

---
 Sources/SwiftCompilerPlugin/CompilerPlugin.swift               | 2 +-
 .../CompilerPluginMessageHandler.swift                         | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/Sources/SwiftCompilerPlugin/CompilerPlugin.swift b/Sources/SwiftCompilerPlugin/CompilerPlugin.swift
index 3ced3a43de1..b9c9073d154 100644
--- a/Sources/SwiftCompilerPlugin/CompilerPlugin.swift
+++ b/Sources/SwiftCompilerPlugin/CompilerPlugin.swift
@@ -112,6 +112,6 @@ extension CompilerPlugin {
     let connection = try StandardIOMessageConnection()
     let provider = MacroProviderAdapter(plugin: Self())
     let impl = CompilerPluginMessageListener(connection: connection, provider: provider)
-    impl.main()
+    try impl.main()
   }
 }
diff --git a/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift b/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift
index da680bb2e20..97e9e3e5401 100644
--- a/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift
+++ b/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift
@@ -96,12 +96,13 @@ public class CompilerPluginMessageListener<Connection: MessageConnection, Handle
   ///
   /// On internal errors, such as I/O errors or JSON serialization errors, print
   /// an error message and `exit(1)`
-  public func main() {
+  public func main() throws {
     #if os(WASI)
     // Rather than blocking on read(), let the host tell us when there's data.
     readabilityHandler = { _ = self.handleNextMessage() }
     #else
     while handleNextMessage() {}
+    try self.handler.shutDown()
     #endif
   }