feat: impl kcl plugin runtime feature gate.

kcl-lang · Sep 12, 2023 · 51f534f · 51f534f
1 parent 8a9c8fa
commit 51f534f
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 11 deletions.
diff --git a/pkg/kclvm_runtime/init.go b/pkg/kclvm_runtime/init.go
@@ -4,16 +4,19 @@ package kclvm_runtime
 
 import (
 	"fmt"
+	"os"
+	"path/filepath"
 	"runtime"
 	"sync"
 
 	"kcl-lang.io/kcl-go/pkg/spec/gpyrpc"
 )
 
 var (
-	Debug        bool
-	pyrpcRuntime *Runtime
-	once         sync.Once
+	Debug              bool
+	rpcRuntime         *Runtime
+	once               sync.Once
+	UseKCLPluginEnvVar = "KCL_GO_USE_PLUGIN"
 )
 
 const tip = "Tip: Have you used a binary version of KCL in your PATH that is not consistent with the KCL Go SDK? You can upgrade or reduce the KCL version or delete the KCL in your PATH"
@@ -24,12 +27,12 @@ func InitRuntime(maxProc int) {
 
 func GetRuntime() *Runtime {
 	once.Do(func() { initRuntime(0) })
-	return pyrpcRuntime
+	return rpcRuntime
 }
 
 func GetPyRuntime() *Runtime {
 	once.Do(func() { initRuntime(0) })
-	return pyrpcRuntime
+	return rpcRuntime
 }
 
 func initRuntime(maxProc int) {
@@ -44,11 +47,18 @@ func initRuntime(maxProc int) {
 		panic(ErrKclvmRootNotFound)
 	}
 
-	pyrpcRuntime = NewRuntime(int(maxProc), "kclvm_cli", "server")
-	pyrpcRuntime.Start()
+	if os.Getenv(UseKCLPluginEnvVar) != "" {
+		os.Setenv("PYTHONHOME", "")
+		os.Setenv("PYTHONPATH", filepath.Join(g_KclvmRoot, "lib", "site-packages"))
+		rpcRuntime = NewRuntime(int(maxProc), MustGetKclvmPath(), "-m", "kclvm.program.rpc-server")
+	} else {
+		rpcRuntime = NewRuntime(int(maxProc), "kclvm_cli", "server")
+	}
+
+	rpcRuntime.Start()
 
 	client := &BuiltinServiceClient{
-		Runtime: pyrpcRuntime,
+		Runtime: rpcRuntime,
 	}
 
 	// ping

diff --git a/pkg/kclvm_runtime/kclvm.go b/pkg/kclvm_runtime/kclvm.go
@@ -23,7 +23,7 @@ const (
 )
 
 func init() {
-	if os.Getenv(DisableArtifactEnvVar) == "" {
+	if os.Getenv(DisableArtifactEnvVar) == "" && os.Getenv(UseKCLPluginEnvVar) == "" {
 		installKclArtifact()
 	}
 	g_KclvmRoot = findKclvmRoot()

diff --git a/pkg/kclvm_runtime/runtime.go b/pkg/kclvm_runtime/runtime.go
@@ -49,9 +49,11 @@ func (p *Runtime) Start() {
 
 	for i, proc := range p.procs {
 		if proc == nil || proc.IsExited() {
-			if proc, err := createProcess(p.exe, p.args...); err == nil {
-				p.procs[i] = proc
+			proc, err := createProcess(p.exe, p.args...)
+			if err != nil {
+				panic(err)
 			}
+			p.procs[i] = proc
 		}
 	}
 }