From 48cab9c77c30fb0b766bc175a7036e43bdf0db81 Mon Sep 17 00:00:00 2001
From: Abhinav Gupta <abhinav@pulumi.com>
Date: Mon, 28 Aug 2023 14:37:23 -0700
Subject: [PATCH] plugin(language, converter): Respond to SIGINT

Configures the language plugin binary to shut down the plugin RPC server
when SIGINT is received.
This will help the binary exit cleanly when it's time to shut down.

This functionality was added in pulumi/pulumi#13795.
pulumi/pulumi#13809 contains a similar change for Go, Python, and Node.
---
 CHANGELOG_PENDING.md              |  4 +++-
 cmd/pulumi-converter-yaml/main.go | 11 +++++++++++
 cmd/pulumi-language-yaml/main.go  | 11 +++++++++--
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG_PENDING.md b/CHANGELOG_PENDING.md
index f32d4b26..d0fe25f9 100644
--- a/CHANGELOG_PENDING.md
+++ b/CHANGELOG_PENDING.md
@@ -5,4 +5,6 @@
 
 - Publish pulumi-converter-yaml.
 
-### Bug Fixes
\ No newline at end of file
+- Plugins: clean up resources and exit cleanly on receiving SIGINT or CTRL_BREAK.
+
+### Bug Fixes
diff --git a/cmd/pulumi-converter-yaml/main.go b/cmd/pulumi-converter-yaml/main.go
index b402ec23..8cfafc4b 100644
--- a/cmd/pulumi-converter-yaml/main.go
+++ b/cmd/pulumi-converter-yaml/main.go
@@ -21,6 +21,8 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"os"
+	"os/signal"
 
 	yamlgen "github.com/pulumi/pulumi-yaml/pkg/pulumiyaml/codegen"
 	"github.com/pulumi/pulumi/pkg/v3/codegen/pcl"
@@ -94,8 +96,17 @@ func (*yamlConverter) ConvertProgram(ctx context.Context,
 
 // Launches the converter RPC endpoint
 func main() {
+	cancelch := make(chan bool)
+	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
+	go func() {
+		<-ctx.Done()
+		cancel() // deregister the interrupt handler
+		close(cancelch)
+	}()
+
 	// Fire up a gRPC server, letting the kernel choose a free port for us.
 	handle, err := rpcutil.ServeWithOptions(rpcutil.ServeOptions{
+		Cancel: cancelch,
 		Init: func(srv *grpc.Server) error {
 			pulumirpc.RegisterConverterServer(srv, plugin.NewConverterServer(&yamlConverter{}))
 			return nil
diff --git a/cmd/pulumi-language-yaml/main.go b/cmd/pulumi-language-yaml/main.go
index c11cb063..482ea61c 100644
--- a/cmd/pulumi-language-yaml/main.go
+++ b/cmd/pulumi-language-yaml/main.go
@@ -20,6 +20,8 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"os"
+	"os/signal"
 	"time"
 
 	"github.com/pkg/errors"
@@ -81,13 +83,18 @@ func main() {
 }
 
 func setupHealthChecks(engineAddress string) (chan bool, error) {
-	// If we have a host cancel our cancellation context if it fails the healthcheck
-	ctx, cancel := context.WithCancel(context.Background())
+	// If the health check begins failing or we receive a SIGINT,
+	// we'll cancel the context.
+	//
+	// The returned channel is used to notify the server that it should
+	// stop serving and exit.
+	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
 
 	// map the context Done channel to the rpcutil boolean cancel channel
 	cancelChannel := make(chan bool)
 	go func() {
 		<-ctx.Done()
+		cancel() // deregister the signal handler
 		close(cancelChannel)
 	}()
 	err := rpcutil.Healthcheck(ctx, engineAddress, 5*time.Minute, cancel)