From efecf0a4f87facacbcaa217bdd05039ec2dc8cfa Mon Sep 17 00:00:00 2001
From: peefy <xpf6677@163.com>
Date: Wed, 24 Jul 2024 13:25:39 +0800
Subject: [PATCH] fix: runtime init panic

Signed-off-by: peefy <xpf6677@163.com>
---
 pkg/runtime/init.go      | 29 ++---------------------------
 pkg/runtime/init_test.go | 28 ++++++++++++++++++++++++++++
 pkg/runtime/kclvm.go     | 15 +++++++++++----
 3 files changed, 41 insertions(+), 31 deletions(-)
 create mode 100644 pkg/runtime/init_test.go

diff --git a/pkg/runtime/init.go b/pkg/runtime/init.go
index 0414c2b1..07ac5ca3 100644
--- a/pkg/runtime/init.go
+++ b/pkg/runtime/init.go
@@ -3,18 +3,10 @@
 package runtime
 
 import (
-	"context"
 	"fmt"
-	"os"
-	"path/filepath"
 	"runtime"
 	"sync"
-	"time"
 
-	"github.com/gofrs/flock"
-	"kcl-lang.io/kcl-go/pkg/env"
-	"kcl-lang.io/kcl-go/pkg/logger"
-	"kcl-lang.io/kcl-go/pkg/path"
 	"kcl-lang.io/kcl-go/pkg/spec/gpyrpc"
 )
 
@@ -36,25 +28,8 @@ func GetRuntime() *Runtime {
 }
 
 func initRuntime(maxProc int) {
-	if !env.GetDisableInstallArtifact() {
-		// Get the install lib path.
-		path := path.LibPath()
-		err := os.MkdirAll(path, 0777)
-		if err != nil {
-			logger.GetLogger().Warningf("install kclvm failed: %s", err.Error())
-		}
-		// Acquire a file lock for process synchronization
-		lockPath := filepath.Join(path, "init.lock")
-		fileLock := flock.New(lockPath)
-		lockCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
-		defer cancel()
-		locked, err := fileLock.TryLockContext(lockCtx, time.Second)
-		if err == nil && locked {
-			defer fileLock.Unlock()
-		}
-	}
 	if maxProc <= 0 {
-		maxProc = 1
+		maxProc = 2
 	}
 	if maxProc > runtime.NumCPU()*2 {
 		maxProc = runtime.NumCPU() * 2
@@ -76,7 +51,7 @@ func initRuntime(maxProc int) {
 		args := &gpyrpc.Ping_Args{Value: "echo"}
 		resp, err := client.Ping(args)
 		if err != nil || resp.Value != args.Value {
-			fmt.Println("Init kcl runtime failed, path: ", MustGetKclvmPath())
+			fmt.Println("Init kcl runtime failed, path:", MustGetKclvmPath())
 			fmt.Println(tip)
 			fmt.Printf("If not, you can run `rm -r %s/bin` to fix this issue\n", MustGetKclvmPath())
 			panic(err)
diff --git a/pkg/runtime/init_test.go b/pkg/runtime/init_test.go
new file mode 100644
index 00000000..a84d4223
--- /dev/null
+++ b/pkg/runtime/init_test.go
@@ -0,0 +1,28 @@
+// Copyright The KCL Authors. All rights reserved.
+
+package runtime
+
+import (
+	"sync"
+	"testing"
+)
+
+func TestInitRuntimeConcurrency(t *testing.T) {
+	var wg sync.WaitGroup
+	numConcurrentCalls := 10
+
+	for i := 0; i < numConcurrentCalls; i++ {
+		wg.Add(1)
+		go func(maxProc int) {
+			defer wg.Done()
+			defer func() {
+				if err := recover(); err != nil {
+					t.Errorf("initRuntime panicked: %v", err)
+				}
+			}()
+			initRuntime(maxProc)
+		}(2)
+	}
+
+	wg.Wait()
+}
diff --git a/pkg/runtime/kclvm.go b/pkg/runtime/kclvm.go
index b9cfbc06..330249e7 100644
--- a/pkg/runtime/kclvm.go
+++ b/pkg/runtime/kclvm.go
@@ -39,12 +39,19 @@ func installKclArtifact() {
 	lockCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 	locked, err := fileLock.TryLockContext(lockCtx, time.Second)
-	if err == nil && locked {
-		defer fileLock.Unlock()
-	}
 	if err != nil {
-		logger.GetLogger().Warningf("install kclvm failed: %s", err.Error())
+		logger.GetLogger().Warningf("attempt to acquire file lock failed: %s", err.Error())
+		return
 	}
+	if !locked {
+		logger.GetLogger().Warning("unable to acquire file lock")
+		return
+	}
+	defer func() {
+		if err := fileLock.Unlock(); err != nil {
+			logger.GetLogger().Errorf("failed to unlock file: %s", err.Error())
+		}
+	}()
 	// Install lib
 	err = lib.InstallKclvm(path)
 	if err != nil {