feat: add support for networked provisioners (#9593)

* Refactor provisionerd to use interface to connect to provisioners

Signed-off-by: Spike Curtis <spike@coder.com>

* feat: add support for networked provisioners

Signed-off-by: Spike Curtis <spike@coder.com>

* fix token length and linting

Signed-off-by: Spike Curtis <spike@coder.com>

---------

Signed-off-by: Spike Curtis <spike@coder.com>

provisionersdk/serve.go

@@ -10,6 +10,7 @@ import (
 	"github.com/hashicorp/yamux"
 	"github.com/valyala/fasthttp/fasthttputil"
 	"golang.org/x/xerrors"
+	"storj.io/drpc"
 	"storj.io/drpc/drpcmux"
 	"storj.io/drpc/drpcserver"
 
@@ -21,8 +22,10 @@ import (
 
 // ServeOptions are configurations to serve a provisioner.
 type ServeOptions struct {
-	// Conn specifies a custom transport to serve the dRPC connection.
-	Listener      net.Listener
+	// Listener serves multiple connections. Cannot be combined with Conn.
+	Listener net.Listener
+	// Conn is a single connection to serve. Cannot be combined with Listener.
+	Conn          drpc.Transport
 	Logger        slog.Logger
 	WorkDirectory string
 }
@@ -38,8 +41,11 @@ func Serve(ctx context.Context, server Server, options *ServeOptions) error {
 	if options == nil {
 		options = &ServeOptions{}
 	}
-	// Default to using stdio.
-	if options.Listener == nil {
+	if options.Listener != nil && options.Conn != nil {
+		return xerrors.New("specify Listener or Conn, not both")
+	}
+	// Default to using stdio with yamux as a Listener
+	if options.Listener == nil && options.Conn == nil {
 		config := yamux.DefaultConfig()
 		config.LogOutput = io.Discard
 		stdio, err := yamux.Server(&readWriteCloser{
@@ -75,10 +81,12 @@ func Serve(ctx context.Context, server Server, options *ServeOptions) error {
 		return xerrors.Errorf("register provisioner: %w", err)
 	}
 	srv := drpcserver.New(&tracing.DRPCHandler{Handler: mux})
-	// Only serve a single connection on the transport.
-	// Transports are not multiplexed, and provisioners are
-	// short-lived processes that can be executed concurrently.
-	err = srv.Serve(ctx, options.Listener)
+
+	if options.Listener != nil {
+		err = srv.Serve(ctx, options.Listener)
+	} else if options.Conn != nil {
+		err = srv.ServeOne(ctx, options.Conn)
+	}
 	if err != nil {
 		if errors.Is(err, io.EOF) ||
 			errors.Is(err, context.Canceled) ||

provisionersdk/serve_test.go

@@ -2,14 +2,17 @@ package provisionersdk_test
 
 import (
 	"context"
+	"net"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
+	"storj.io/drpc/drpcconn"
 
 	"github.com/coder/coder/v2/provisionersdk"
 	"github.com/coder/coder/v2/provisionersdk/proto"
+	"github.com/coder/coder/v2/testutil"
 )
 
 func TestMain(m *testing.M) {
@@ -18,7 +21,7 @@ func TestMain(m *testing.M) {
 
 func TestProvisionerSDK(t *testing.T) {
 	t.Parallel()
-	t.Run("Serve", func(t *testing.T) {
+	t.Run("ServeListener", func(t *testing.T) {
 		t.Parallel()
 		client, server := provisionersdk.MemTransportPipe()
 		defer client.Close()
@@ -72,6 +75,61 @@ func TestProvisionerSDK(t *testing.T) {
 		})
 		require.NoError(t, err)
 	})
+
+	t.Run("ServeConn", func(t *testing.T) {
+		t.Parallel()
+		client, server := net.Pipe()
+		defer client.Close()
+		defer server.Close()
+
+		ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitMedium)
+		defer cancelFunc()
+		srvErr := make(chan error, 1)
+		go func() {
+			err := provisionersdk.Serve(ctx, unimplementedServer{}, &provisionersdk.ServeOptions{
+				Conn:          server,
+				WorkDirectory: t.TempDir(),
+			})
+			srvErr <- err
+		}()
+
+		api := proto.NewDRPCProvisionerClient(drpcconn.New(client))
+		s, err := api.Session(ctx)
+		require.NoError(t, err)
+		err = s.Send(&proto.Request{Type: &proto.Request_Config{Config: &proto.Config{}}})
+		require.NoError(t, err)
+
+		err = s.Send(&proto.Request{Type: &proto.Request_Parse{Parse: &proto.ParseRequest{}}})
+		require.NoError(t, err)
+		msg, err := s.Recv()
+		require.NoError(t, err)
+		require.Equal(t, "unimplemented", msg.GetParse().GetError())
+
+		err = s.Send(&proto.Request{Type: &proto.Request_Plan{Plan: &proto.PlanRequest{}}})
+		require.NoError(t, err)
+		msg, err = s.Recv()
+		require.NoError(t, err)
+		// Plan has no error so that we're allowed to run Apply
+		require.Equal(t, "", msg.GetPlan().GetError())
+
+		err = s.Send(&proto.Request{Type: &proto.Request_Apply{Apply: &proto.ApplyRequest{}}})
+		require.NoError(t, err)
+		msg, err = s.Recv()
+		require.NoError(t, err)
+		require.Equal(t, "unimplemented", msg.GetApply().GetError())
+
+		// Check provisioner closes when the connection does
+		err = s.Close()
+		require.NoError(t, err)
+		err = api.DRPCConn().Close()
+		require.NoError(t, err)
+		select {
+		case <-ctx.Done():
+			t.Fatal("timeout waiting for provisioner")
+		case err = <-srvErr:
+			require.NoError(t, err)
+		}
+	})
 }
 
 type unimplementedServer struct{}