plumb vm getter to runtimetype

hexfusion · Jun 23, 2023 · 2455465 · 2455465
1 parent eb52fee
commit 2455465
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 42 deletions.
diff --git a/vms/registry/vm_getter.go b/vms/registry/vm_getter.go
@@ -112,16 +112,21 @@ func (getter *vmGetter) Get() (map[ids.ID]vms.Factory, map[ids.ID]vms.Factory, e
 			return nil, nil, err
 		}
 
+		var runtimeType rpcchainvm.RuntimeType
 		if !isPod {
-			unregisteredVMs[vmID] = rpcchainvm.NewFactory(
-				filepath.Join(getter.config.PluginDirectory, file.Name()),
-				getter.config.CPUTracker,
-				getter.config.RuntimeTracker,
-			)
+			runtimeType = rpcchainvm.RuntimeSubprocess
 		} else {
 			// TODO: use the container implementation here
 			fmt.Printf("Attempt to use the container implementation to create a VM Factory for: %s\n", file.Name())
+			runtimeType = rpcchainvm.RuntimeContainer
 		}
+
+		unregisteredVMs[vmID] = rpcchainvm.NewFactory(
+			filepath.Join(getter.config.PluginDirectory, file.Name()),
+			runtimeType,
+			getter.config.CPUTracker,
+			getter.config.RuntimeTracker,
+		)
 	}
 	return registeredVMs, unregisteredVMs, nil
 }
diff --git a/vms/rpcchainvm/factory.go b/vms/rpcchainvm/factory.go
@@ -6,65 +6,112 @@ package rpcchainvm
 import (
 	"context"
 	"fmt"
+	"os"
 
 	"github.com/ava-labs/avalanchego/utils/logging"
 	"github.com/ava-labs/avalanchego/utils/resource"
 	"github.com/ava-labs/avalanchego/vms"
 	"github.com/ava-labs/avalanchego/vms/rpcchainvm/grpcutils"
 	"github.com/ava-labs/avalanchego/vms/rpcchainvm/runtime"
+	"github.com/ava-labs/avalanchego/vms/rpcchainvm/runtime/container"
 	"github.com/ava-labs/avalanchego/vms/rpcchainvm/runtime/subprocess"
 
 	vmpb "github.com/ava-labs/avalanchego/proto/pb/vm"
 )
 
 var _ vms.Factory = (*factory)(nil)
 
+type RuntimeType string
+
+const (
+	RuntimeContainer  RuntimeType = "Container"
+	RuntimeSubprocess RuntimeType = "Subprocess"
+)
+
 type factory struct {
+	runtime        RuntimeType
 	path           string
 	processTracker resource.ProcessTracker
 	runtimeTracker runtime.Tracker
 }
 
-func NewFactory(path string, processTracker resource.ProcessTracker, runtimeTracker runtime.Tracker) vms.Factory {
+func NewFactory(path string, runtimeType RuntimeType, processTracker resource.ProcessTracker, runtimeTracker runtime.Tracker) vms.Factory {
 	return &factory{
+		runtime:        runtimeType,
 		path:           path,
 		processTracker: processTracker,
 		runtimeTracker: runtimeTracker,
 	}
 }
 
 func (f *factory) New(log logging.Logger) (interface{}, error) {
-	config := &subprocess.Config{
-		Stderr:           log,
-		Stdout:           log,
-		HandshakeTimeout: runtime.DefaultHandshakeTimeout,
-		Log:              log,
-	}
 
 	listener, err := grpcutils.NewListener()
 	if err != nil {
 		return nil, fmt.Errorf("failed to create listener: %w", err)
 	}
 
-	status, stopper, err := subprocess.Bootstrap(
-		context.TODO(),
-		listener,
-		subprocess.NewCmd(f.path),
-		config,
-	)
-	if err != nil {
-		return nil, err
-	}
+	var vm *VMClient
 
-	clientConn, err := grpcutils.Dial(status.Addr)
-	if err != nil {
-		return nil, err
-	}
+	// TODO dedupe O:)
+	switch f.runtime {
+	case RuntimeSubprocess:
+		config := &subprocess.Config{
+			Stderr:           log,
+			Stdout:           log,
+			HandshakeTimeout: runtime.DefaultHandshakeTimeout,
+			Log:              log,
+		}
+
+		status, stopper, err := subprocess.Bootstrap(
+			context.TODO(),
+			listener,
+			subprocess.NewCmd(f.path),
+			config,
+		)
+		if err != nil {
+			return nil, err
+		}
+
+		clientConn, err := grpcutils.Dial(status.Addr)
+		if err != nil {
+			return nil, err
+		}
 
-	vm := NewClient(vmpb.NewVMClient(clientConn))
-	vm.SetProcess(stopper, status.Pid, f.processTracker)
+		vm = NewClient(vmpb.NewVMClient(clientConn))
+		vm.SetProcess(stopper, status.Pid, f.processTracker)
+		f.runtimeTracker.TrackRuntime(stopper)
 
-	f.runtimeTracker.TrackRuntime(stopper)
+	case RuntimeContainer:
+		podPath := fmt.Sprintf("%s_pod.yaml", f.path)
+		podBytes, err := os.ReadFile(podPath)
+		if err != nil {
+			return nil, fmt.Errorf("failed to read pod yaml: %q", podPath)
+		}
+		config := &container.Config{
+			PodBytes:         podBytes,
+			HandshakeTimeout: runtime.DefaultHandshakeTimeout,
+			Log:              log,
+		}
+
+		status, stopper, err := container.Bootstrap(
+			context.TODO(),
+			listener,
+			config,
+		)
+		if err != nil {
+			return nil, err
+		}
+		clientConn, err := grpcutils.Dial(status.Addr)
+		if err != nil {
+			return nil, err
+		}
+
+		vm = NewClient(vmpb.NewVMClient(clientConn))
+		// IDK if this does anything crazy with pid 0
+		vm.SetProcess(stopper, 0, f.processTracker)
+		f.runtimeTracker.TrackRuntime(stopper)
+	}
 
 	return vm, nil
 }
diff --git a/vms/rpcchainvm/runtime/container/client.go b/vms/rpcchainvm/runtime/container/client.go
@@ -76,4 +76,4 @@ func getSocketPath() (string, error) {
 		return "", fmt.Errorf("failed to find rootless socket")
 	}
 	return fmt.Sprintf("unix:%s/podman/podman.sock", sockDir), nil
-}
+}
diff --git a/vms/rpcchainvm/runtime/container/runtime.go b/vms/rpcchainvm/runtime/container/runtime.go
@@ -7,17 +7,13 @@ import (
 	"bytes"
 	"context"
 	"fmt"
-	"io"
 	"net"
-	"os"
-	"strings"
 	"time"
 
+	"github.com/ghodss/yaml"
 	"go.uber.org/zap"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes/scheme"
-	"github.com/ghodss/yaml"
-
 
 	"github.com/ava-labs/avalanchego/utils/logging"
 	"github.com/ava-labs/avalanchego/vms/rpcchainvm/grpcutils"
@@ -30,10 +26,7 @@ import (
 )
 
 type Config struct {
-	// Stderr of the VM process written to this writer.
-	Stderr io.Writer
-	// Stdout of the VM process written to this writer.
-	Stdout io.Writer
+	PodBytes []byte
 	// Duration engine server will wait for handshake success.
 	HandshakeTimeout time.Duration
 	Log              logging.Logger
@@ -83,7 +76,7 @@ func Bootstrap(
 	}
 
 	// all this should go into factory
-	obj, _, err := scheme.Codecs.UniversalDeserializer().Decode([]byte(podYamlBytes), nil, nil)
+	obj, _, err := scheme.Codecs.UniversalDeserializer().Decode(config.PodBytes, nil, nil)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to derialize pod yaml: %w", err)
 	}
@@ -136,8 +129,8 @@ func Bootstrap(
 	)
 
 	status := &Status{
-		PodBytes:  podBytes,
-		Addr: intitializer.vmAddr,
+		PodBytes: podBytes,
+		Addr:     intitializer.vmAddr,
 	}
 	return status, stopper, nil
 }
diff --git a/vms/rpcchainvm/runtime/container/stopper.go b/vms/rpcchainvm/runtime/container/stopper.go
@@ -34,7 +34,6 @@ func (s *stopper) Stop(ctx context.Context) {
 	})
 }
 
-
 func stop(ctx context.Context, log logging.Logger, cmd *exec.Cmd) {
 	waitChan := make(chan error)
 	go func() {