refactor: optimize logs (#268)

hengyoush · Jan 12, 2025 · 57c7542 · 57c7542
1 parent 8b17632
commit 57c7542
Show file tree

Hide file tree

Showing 5 changed files with 244 additions and 131 deletions.
diff --git a/agent/agent.go b/agent/agent.go
@@ -98,13 +98,14 @@ func SetupAgent(options ac.AgentOptions) {
 
 	var _bf loader.BPF
 	go func(_bf *loader.BPF) {
+		defer wg.Done()
 		options.LoadPorgressChannel <- "🍩 Kyanos starting..."
 		kernelVersion := compatible.GetCurrentKernelVersion()
 		options.Kv = &kernelVersion
 		var err error
 		defer func() {
 			if err != nil {
-				common.AgentLog.Error("Failed to load BPF programs: ", err)
+				common.AgentLog.Errorf("Failed to load BPF programs: %+v", errors.Unwrap(errors.Unwrap(err)))
 				_bf.Err = err
 				options.LoadPorgressChannel <- "❌ Kyanos start failed"
 				options.LoadPorgressChannel <- "quit"
@@ -151,7 +152,6 @@ func SetupAgent(options ac.AgentOptions) {
 			time.Sleep(500 * time.Millisecond)
 			options.LoadPorgressChannel <- "quit"
 		}
-		defer wg.Done()
 	}(&_bf)
 	defer func() {
 		_bf.Close()

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -586,25 +586,25 @@ func TestSslEventsCanRelatedToKernEvents(t *testing.T) {
 			bpf.AttachSyscallWriteExit,
 			bpf.AttachKProbeSecuritySocketSendmsgEntry,
 			bpf.AttachKProbeSecuritySocketRecvmsgEntry,
-			func() link.Link {
+			func() (link.Link, error) {
 				return ApplyKernelVersionFunctions(t, bpf.AgentStepTDEV_IN)
 			},
-			func() link.Link {
+			func() (link.Link, error) {
 				return ApplyKernelVersionFunctions(t, bpf.AgentStepTUSER_COPY)
 			},
-			func() link.Link {
+			func() (link.Link, error) {
 				return ApplyKernelVersionFunctions(t, bpf.AgentStepTIP_IN)
 			},
-			func() link.Link {
+			func() (link.Link, error) {
 				return ApplyKernelVersionFunctions(t, bpf.AgentStepTTCP_IN)
 			},
-			func() link.Link {
+			func() (link.Link, error) {
 				return ApplyKernelVersionFunctions(t, bpf.AgentStepTIP_OUT)
 			},
-			func() link.Link {
+			func() (link.Link, error) {
 				return ApplyKernelVersionFunctions(t, bpf.AgentStepTDEV_OUT)
 			},
-			func() link.Link {
+			func() (link.Link, error) {
 				return ApplyKernelVersionFunctions(t, bpf.AgentStepTQDISC_OUT)
 			},
 		},
@@ -964,7 +964,7 @@ func TestIpXmit(t *testing.T) {
 			bpf.AttachSyscallWriteEntry,
 			bpf.AttachSyscallWriteExit,
 			bpf.AttachKProbeSecuritySocketSendmsgEntry,
-			func() link.Link {
+			func() (link.Link, error) {
 				return ApplyKernelVersionFunctions(t, bpf.AgentStepTIP_OUT)
 			},
 		},
@@ -994,10 +994,10 @@ func TestDevQueueXmit(t *testing.T) {
 		bpf.AttachSyscallWriteEntry,
 		bpf.AttachSyscallWriteExit,
 		bpf.AttachKProbeSecuritySocketSendmsgEntry,
-		func() link.Link {
+		func() (link.Link, error) {
 			return ApplyKernelVersionFunctions(t, bpf.AgentStepTIP_OUT)
 		},
-		func() link.Link {
+		func() (link.Link, error) {
 			return ApplyKernelVersionFunctions(t, bpf.AgentStepTQDISC_OUT)
 		},
 	}, "GET DevQueueXmit\n", Write, Read,
@@ -1029,10 +1029,10 @@ func TestDevHardStartXmit(t *testing.T) {
 			bpf.AttachSyscallWriteEntry,
 			bpf.AttachSyscallWriteExit,
 			bpf.AttachKProbeSecuritySocketSendmsgEntry,
-			func() link.Link {
+			func() (link.Link, error) {
 				return ApplyKernelVersionFunctions(t, bpf.AgentStepTIP_OUT)
 			},
-			func() link.Link {
+			func() (link.Link, error) {
 				return ApplyKernelVersionFunctions(t, bpf.AgentStepTDEV_OUT)
 			},
 		}, "GET DevHardStartXmit\n", Write, Read,
@@ -1110,10 +1110,10 @@ func TestIpRcvCore(t *testing.T) {
 		bpf.AttachSyscallWriteExit,
 		bpf.AttachKProbeSecuritySocketSendmsgEntry,
 		bpf.AttachKProbeSecuritySocketRecvmsgEntry,
-		func() link.Link {
+		func() (link.Link, error) {
 			return ApplyKernelVersionFunctions(t, bpf.AgentStepTDEV_IN)
 		},
-		func() link.Link {
+		func() (link.Link, error) {
 			return ApplyKernelVersionFunctions(t, bpf.AgentStepTIP_IN)
 		},
 	},
@@ -1150,10 +1150,10 @@ func TestTcpV4DoRcv(t *testing.T) {
 		bpf.AttachSyscallWriteExit,
 		bpf.AttachKProbeSecuritySocketSendmsgEntry,
 		bpf.AttachKProbeSecuritySocketRecvmsgEntry,
-		func() link.Link {
+		func() (link.Link, error) {
 			return ApplyKernelVersionFunctions(t, bpf.AgentStepTDEV_IN)
 		},
-		func() link.Link {
+		func() (link.Link, error) {
 			return ApplyKernelVersionFunctions(t, bpf.AgentStepTTCP_IN)
 		},
 	},
@@ -1190,10 +1190,10 @@ func TestSkbCopyDatagramIter(t *testing.T) {
 		bpf.AttachSyscallWriteExit,
 		bpf.AttachKProbeSecuritySocketSendmsgEntry,
 		bpf.AttachKProbeSecuritySocketRecvmsgEntry,
-		func() link.Link {
+		func() (link.Link, error) {
 			return ApplyKernelVersionFunctions(t, bpf.AgentStepTDEV_IN)
 		},
-		func() link.Link {
+		func() (link.Link, error) {
 			return ApplyKernelVersionFunctions(t, bpf.AgentStepTUSER_COPY)
 		},
 	},

diff --git a/agent/agent_utils_test.go b/agent/agent_utils_test.go
@@ -52,7 +52,11 @@ func StartAgent0(bpfAttachFunctions []bpf.AttachBpfProgFunction,
 			progs := list.New()
 			for _, each := range bpfAttachFunctions {
 				if each != nil {
-					progs.PushBack(each())
+					l, err := each()
+					if err != nil {
+						log.Fatalf("Attach failed: %v", err)
+					}
+					progs.PushBack(l)
 				}
 			}
 			return progs
@@ -777,15 +781,15 @@ func min(a, b int) int {
 // 	compatilbeMode = b
 // }
 
-func ApplyKernelVersionFunctions(t *testing.T, step bpf.AgentStepT) link.Link {
+func ApplyKernelVersionFunctions(t *testing.T, step bpf.AgentStepT) (link.Link, error) {
 	v := compatible.GetCurrentKernelVersion()
 	if step == bpf.AgentStepTNIC_IN {
 		if v.SupportCapability(compatible.SupportXDP) {
 			l, err := bpf.AttachXdp()
 			if err != nil {
 				t.Fatal(err)
 			} else {
-				return l
+				return l, nil
 			}
 		} else {
 			t.FailNow()
@@ -813,11 +817,11 @@ func ApplyKernelVersionFunctions(t *testing.T, step bpf.AgentStepT) link.Link {
 				log.Fatalf("Attach failed: %v, functions: %v", err, functions)
 			}
 		} else {
-			return l
+			return l, nil
 		}
 	}
 	t.FailNow()
-	return nil
+	return nil, nil
 }
 
 func KernRcvTestWithHTTP(t *testing.T, progs []bpf.AttachBpfProgFunction, kernEvtFilter FindInterestedKernEventOptions, kernEvtAsserts KernDataEventAssertConditions) {