fix pod添加超时控制，避免因容器卡死而卡死

cmd/pid_of_docker/main.go

@@ -11,11 +11,13 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"time"
 
 	"get-container/cmd/pid_of_docker/logic"
 
 	"github.com/charmbracelet/lipgloss"
 	"github.com/moby/moby/client"
+	"github.com/spf13/pflag"
 )
 
 /*
@@ -57,15 +59,30 @@ func (c *Cinfo) Format() ([5]string, [5]int) {
 var (
 	RegUser       = regexp.MustCompile(`^(?i)/public[0-9]*/home/([0-9a-z_]+)(?:|/.*)$`)
 	RegUserPublic = regexp.MustCompile(`^(?i)/public[0-9]*/home/(?:public_user|locauser|localuser)/([0-9a-z_]+)(?:|/.*)$`)
+
+	flagHelp = pflag.BoolP("help", "h", false, "show usage")
 )
 
 func main() {
+
+	pflag.Parse()
+	if *flagHelp {
+		fmt.Println("这个工具用于查看docker容器的创建者和使用计算卡的情况")
+		fmt.Println("容器创建者由容器挂载的目录和标签com.sugon.username推导出来")
+		fmt.Println("本工具支持查看nvidia和hycu")
+		os.Exit(0)
+	}
+
 	cli, err := docker.GetDockerClient()
 	if err != nil {
 		log.Fatalf("can't connect to docker daemon: %v", err)
 	}
 
-	csum, err := cli.ContainerList(context.Background(), client.ContainerListOptions{})
+	// docker ps 超时控制
+	ctx, canfunc := context.WithTimeout(context.Background(), time.Second)
+	defer canfunc()
+
+	csum, err := cli.ContainerList(ctx, client.ContainerListOptions{})
 	if err != nil {
 		log.Printf("error get container list: %v \n", err)
 		cli.Close()
@@ -99,13 +116,18 @@ func main() {
 				}
 			}
 		}
-
-		top, err := cli.ContainerTop(context.Background(), i.ID, nil)
+		// docker top 超时控制
+		ctxTout, ctxFunc := context.WithTimeout(context.Background(), time.Second)
+		top, err := cli.ContainerTop(ctxTout, i.ID, nil)
 		if err != nil {
+			ctxFunc()
+			fmt.Printf("get pid of container %s timeout: %v \n", i.ID, err)
 			continue
 		}
+		ctxFunc()
 		index := slices.Index(top.Titles, "PID")
 		if index == -1 {
+			fmt.Printf("cat't find PID field in ContainerID: %s \n", i.ID)
 			continue
 		}
 		c.Pid = make([]int, 0, len(top.Processes))
@@ -115,7 +137,6 @@ func main() {
 				c.Pid = append(c.Pid, pid)
 			}
 		}
-
 	}
 	cli.Close()
 

docker/find_test.go

@@ -141,3 +141,27 @@ func TestTop(t *testing.T) {
 	}
 	cli.Close()
 }
+
+func TestTimeout(t *testing.T) {
+	cli, err := GetDockerClient()
+	if err != nil {
+		t.Error(err)
+	}
+	ctx, fun := context.WithTimeout(context.Background(), time.Nanosecond*100)
+	cs, err := cli.ContainerList(ctx, client.ContainerListOptions{})
+	if err != nil {
+		fun()
+		cli.Close()
+		t.Error(err)
+	}
+	fun()
+
+	for _, v := range cs {
+		top, err := cli.ContainerTop(context.Background(), v.ID, nil)
+		if err != nil {
+			continue
+		}
+		t.Logf("%v", top)
+	}
+	cli.Close()
+}