From aa05b04dcc84275534066894334dff290c715de9 Mon Sep 17 00:00:00 2001
From: mrproliu <741550557@qq.com>
Date: Tue, 30 Jul 2024 18:43:26 +0800
Subject: [PATCH] Fix the protocol logs may be missing if the process is
 short-lived (#135)

---
 CHANGES.md                         |  1 +
 pkg/accesslog/common/connection.go | 12 ++++++++----
 pkg/accesslog/runner.go            | 23 +++++++++++++++++++++--
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 9df9eb73..3b22bbae 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -16,6 +16,7 @@ Release Notes.
 * Fix concurrent map operation in the access log module.
 * Fix the profiling cannot found process issue.
 * Fix cannot translate peer address in some UDP scenarios.
+* Fix the protocol logs may be missing if the process is short-lived.
 
 #### Documentation
 
diff --git a/pkg/accesslog/common/connection.go b/pkg/accesslog/common/connection.go
index e339a027..72fbce47 100644
--- a/pkg/accesslog/common/connection.go
+++ b/pkg/accesslog/common/connection.go
@@ -570,7 +570,7 @@ func (c *ConnectionManager) updateMonitorStatusForProcess(pid int32, monitor boo
 func (c *ConnectionManager) OnBuildConnectionLogFinished() {
 	// delete all connections which marked as deletable
 	// all deletable connection events been sent
-	deletableConnections := make([]string, 0)
+	deletableConnections := make(map[string]bool, 0)
 	c.connections.IterCb(func(key string, v interface{}) {
 		con, ok := v.(*ConnectionInfo)
 		if !ok || con == nil {
@@ -580,22 +580,26 @@ func (c *ConnectionManager) OnBuildConnectionLogFinished() {
 		shouldDelete := con.MarkDeletable || !c.ProcessIsMonitor(con.PID)
 
 		if shouldDelete {
-			deletableConnections = append(deletableConnections, key)
+			deletableConnections[key] = true
 		}
 	})
 
 	deleteFromUnfinished := make([]string, 0)
 	for conKey, processorFinished := range c.allUnfinishedConnections {
 		if *processorFinished {
-			deletableConnections = append(deletableConnections, conKey)
+			deletableConnections[conKey] = true
 			deleteFromUnfinished = append(deleteFromUnfinished, conKey)
+		} else {
+			// if the processor not finished, then ignore it from deletable connections
+			delete(deletableConnections, conKey)
 		}
 	}
 	for _, key := range deleteFromUnfinished {
 		delete(c.allUnfinishedConnections, key)
 	}
 
-	for _, key := range deletableConnections {
+	for key := range deletableConnections {
+		log.Debugf("deleting the connection in manager: %s", key)
 		c.connections.Remove(key)
 	}
 }
diff --git a/pkg/accesslog/runner.go b/pkg/accesslog/runner.go
index befac3a6..8164e814 100644
--- a/pkg/accesslog/runner.go
+++ b/pkg/accesslog/runner.go
@@ -22,6 +22,8 @@ import (
 	"fmt"
 	"time"
 
+	process2 "github.com/shirou/gopsutil/process"
+
 	"github.com/sirupsen/logrus"
 
 	"github.com/apache/skywalking-rover/pkg/accesslog/bpf"
@@ -284,6 +286,23 @@ func (r *Runner) convertTimeToInstant(t time.Time) *v32.Instant {
 	}
 }
 
+func (r *Runner) shouldReportProcessLog(pid uint32) bool {
+	// if the process not monitoring, then check the process is existed or not
+	if r.context.ConnectionMgr.ProcessIsMonitor(pid) {
+		return true
+	}
+	exists, err := process2.PidExists(int32(pid))
+	if err != nil {
+		log.Warnf("check pid exists error, pid: %d, error: %v", pid, err)
+		return false
+	}
+	if exists {
+		return false
+	}
+	log.Debugf("the log should be also uploaded because the process quick shutdown but the log exist, pid: %d", pid)
+	return true
+}
+
 func (r *Runner) buildProtocolLog(protocolLog *common.ProtocolLog) (*common.ConnectionInfo,
 	[]*v3.AccessLogKernelLog, *v3.AccessLogProtocolLogs, bool) {
 	if len(protocolLog.KernelLogs) == 0 {
@@ -292,7 +311,7 @@ func (r *Runner) buildProtocolLog(protocolLog *common.ProtocolLog) (*common.Conn
 	firstKernelLog := protocolLog.KernelLogs[0]
 	pid, _ := events.ParseConnectionID(firstKernelLog.GetConnectionID())
 	// if the process not monitoring, then ignore it
-	if !r.context.ConnectionMgr.ProcessIsMonitor(pid) {
+	if !r.shouldReportProcessLog(pid) {
 		return nil, nil, nil, false
 	}
 	connection := r.context.ConnectionMgr.Find(firstKernelLog)
@@ -319,7 +338,7 @@ func (r *Runner) buildProtocolLog(protocolLog *common.ProtocolLog) (*common.Conn
 func (r *Runner) buildKernelLog(kernelLog *common.KernelLog) (*common.ConnectionInfo, *v3.AccessLogKernelLog, bool) {
 	pid, _ := events.ParseConnectionID(kernelLog.Event.GetConnectionID())
 	// if the process not monitoring, then ignore it
-	if !r.context.ConnectionMgr.ProcessIsMonitor(pid) {
+	if !r.shouldReportProcessLog(pid) {
 		return nil, nil, false
 	}
 	connection := r.context.ConnectionMgr.Find(kernelLog.Event)