Filtering process data from mongodb

First of all, fix sth in listener to fit the function of filter. Ori-
ginally, listener mark the /usr/bin/containerd process id with star,
but the children in db is updated by ppid, which is pid of parent but
not tgid, so the stared pid has no children. To Fix this, we add all
the pid of /usr/bin/containerd into the db, and set their ptgid/tgid,
so that they're just normal process as others. Maybe we should finish
the info of these processes? haha.

Then, the filter of pid. There're some designed steps to do, and their
methods are as follows:

- Initially, because of the multithreading execution of listener,
  there may be several entries for the same process, and we should
  merge them. Extract data from database into a slice, and use a map
  to record process info. Iterate the slice, if the pid is in the
  map, then merge them, else insert into the map.
- Then, we should build process tree, but what we have is pid. So
  use another data structure, iterate merged process map, and build
  a map from tgid to a slice of processes. Find out the star. Build
  a map from pid to its tgid.
- BFS. Design a simple queue, and build the tree from the root(stared
  tgid), record all the visited tgid in another map. That's just the
  tree.

As usual, let's talk about the remaining issues:
- Some pids did not recieve exit message. Check the exit time of
  its tgid, or even its ppid.
- Optimize the data structure, record the tree by itself. Now the
  tree is recorded by not only the last helloTree map from tgid to
  slice but the map from pid to tgid. It's hard to store in the
  database. Design a better ds, so the viewer can build the tree
  quickly from the data in db.
- For future file filter, the close time, the same file for the same
  pid, and the pathName of a file, should be paid mych attention.

Fighting!

2 files changed, 275 insertions, 84 deletions

diff --git a/filter/filter.go b/filter/filter.go
index c83fb13..2a774a3 100644
--- a/filter/filter.go
+++ b/filter/filter.go
@@ -5,114 +5,213 @@ import (
         "fmt"
         "log"
         "os"
-        "time"
+        "sync"
 
         "go.mongodb.org/mongo-driver/bson"
         "go.mongodb.org/mongo-driver/mongo"
         "go.mongodb.org/mongo-driver/mongo/options"
-        "go.mongodb.org/mongo-driver/mongo/readpref"
 )
 
-type Exec struct {
-        timestamp time.Time `bson:"timestamp"`
-        execArgs  []string  `bson:"execArgs"`
-}
-
-type Process struct {
-        timestamp      time.Time `bson:"start_timestamp"`
-        ppid           int       `bson:"ppid"`
-        parentTgid     int       `bson:"parentTgid"`
-        pid            int       `bson:"pid"`
-        tgid           int       `bson:"tgid"`
-        args           []string  `bson:"args"`
-        comm           string    `bson:"comm"`
-        cwd            string    `bson:"cwd"`
-        execve         []Exec    `bson:"execve"`
-        exit_code      int       `bson:"exit_code"`
-        exit_signal    int       `bson:"exit_signal"`
-        exit_timestamp time.Time `bson:"exit_timestamp"`
-}
+const (
+        oldDBName     = "test"
+        oldPidColName = "pids"
+)
 
-func (p Process) String() string {
-        var res string
-        res = ""
-        res += fmt.Sprintf("timestamp\t%v\n", p.timestamp)
-        res += fmt.Sprintf("ppid\t%d\nparentTgid\t%d\n", p.ppid, p.parentTgid)
-        res += fmt.Sprintf("pid\t%d\ntgid\t%d\nargs: ", p.pid, p.tgid)
-        for i := 0; i < len(p.args); i++ {
-                res += fmt.Sprintf("%s ", p.args[i])
-        }
-        res += fmt.Sprintf("\ncomm\t%s\ncwd\t%s\n", p.comm, p.cwd)
-        return res
+type treeNode struct {
+        Tgid     int
+        Threads  []Process
+        Children []int
 }
 
-// type Process struct {
-//      StartTimestamp time.Time  `bson:"start_timestamp"`
-//      Ppid           *int       `bson:"ppid"`
-//      ParentTgid     *int       `bson:"parentTgid"`
-//      Pid            int        `bson:"pid"`
-//      Tgid           int        `bson:"tgid"`
-//      Args           []string   `bson:"args"`
-//      Comm           *string    `bson:"comm"`
-//      Cwd            *string    `bson:"cwd"`
-//      Execve         []Exec     `bson:"execve"`
-//      ExitCode       *int       `bson:"exit_code"`
-//      ExitSignal     *int       `bson:"exit_signal"`
-//      ExitTimestamp  *time.Time `bson:"exit_timestamp"`
-// }
-
-// func (p Process) String() string {
-//      var res string
-//      res = ""
-//      res += fmt.Sprintf("timestamp\t%v\n", p.StartTimestamp)
-//      if p.Ppid != nil && p.ParentTgid != nil {
-//              res += fmt.Sprintf("ppid\t%d\nparentTgid\t%d\n", *(p.Ppid), *(p.ParentTgid))
-//      }
-//      res += fmt.Sprintf("pid\t%d\ntgid\t%d\nargs: ", p.Pid, p.Tgid)
-//      for i := 0; i < len(p.Args); i++ {
-//              res += fmt.Sprintf("%s ", p.Args[i])
-//      }
-//      if p.Comm != nil && p.Cwd != nil {
-//              res += fmt.Sprintf("\ncomm\t%s\ncwd\t%s\n", *(p.Comm), *(p.Cwd))
-//      }
-//      return res
-// }
-
 func main() {
+        // 连接到MongoDB
         client, err := mongo.Connect(context.TODO(), options.Client().ApplyURI("mongodb://localhost:27017"))
         if err != nil {
-                fmt.Fprintf(os.Stderr, "Err connecting mongodb: %v\n", err)
+                log.Fatal(err)
         }
         defer client.Disconnect(context.TODO())
 
-        // 检查连接
-        err = client.Ping(context.TODO(), readpref.Primary())
+        // 选择数据库和集合
+        db := client.Database(oldDBName)
+        collection := db.Collection(oldPidColName)
+
+        // 提取所有数据
+        var res []Process
+
+        cursor, err := collection.Find(context.Background(), bson.M{})
         if err != nil {
-                log.Fatal(err)
+                fmt.Fprintf(os.Stderr, "Err: %v\n", err)
+        }
+        if err := cursor.All(context.Background(), &res); err != nil {
+                fmt.Fprintf(os.Stderr, "Err All: %v\n", err)
         }
 
-        pidCol := client.Database("test").Collection("pids")
-        cur, err := pidCol.Find(context.TODO(), bson.M{}) // 查询所有文档
-        if err != nil {
-                log.Fatal(err)
+        var merged sync.Map
+        for _, process := range res {
+                tmp, ok := merged.Load(process.Pid)
+                if ok {
+                        // 证明重复了，要合并
+                        tmp := ProMerge(tmp.(Process), process)
+                        merged.Store(process.Pid, tmp)
+                } else {
+                        // 没有，直接插入
+                        merged.Store(process.Pid, process)
+                }
         }
-        defer cur.Close(context.TODO()) // 确保游标被关闭
 
-        var res []Process
-        for cur.Next(context.TODO()) {
-                var tmp Process
-                // 解码到Process结构体
-                if err := cur.Decode(&tmp); err != nil {
-                        log.Fatal(err)
+        var treeMap sync.Map
+        findTgid := make(map[int]int)
+        var stared int
+        merged.Range(func(key, val interface{}) bool {
+                tmp := val.(Process)
+                if tmp.Star {
+                        stared = tmp.Tgid
+                }
+                // 登记tgid
+                findTgid[tmp.Pid] = tmp.Tgid
+                nodeTmp, ok := treeMap.Load(tmp.Tgid)
+                if ok {
+                        // 直接记录
+                        node := nodeTmp.(treeNode)
+                        node.Threads = append(node.Threads, tmp)
+                        node.Children = append(node.Children, tmp.Children...)
+                        treeMap.Store(tmp.Tgid, node)
+                } else {
+                        node := treeNode{
+                                Tgid:     tmp.Tgid,
+                                Threads:  make([]Process, 0),
+                                Children: make([]int, 0),
+                        }
+                        node.Threads = append(node.Threads, tmp)
+                        node.Children = append(node.Children, tmp.Children...)
+                        treeMap.Store(tmp.Tgid, node)
                 }
-                res = append(res, tmp)
+                return true
+        })
+
+        // 从tgid==stared开始，构建树
+        var helloTree sync.Map // 在树上的tgid节点
+        var q Queue            // 记录每一个整理好的结构体，bfs
+        visited := make(map[int]bool)
+        visited[stared] = true
+        tmp, ok := treeMap.Load(stared)
+        if !ok {
+                return
         }
 
-        if err := cur.Err(); err != nil {
-                log.Fatal(err)
+        q.Enqueue(tmp)
+        helloTree.Store(stared, tmp)
+        for !q.IsEmpty() {
+                tmp, ok := q.Dequeue()
+                if !ok {
+                        continue
+                }
+                node := tmp.(treeNode)
+                for i := 0; i < len(node.Children); i++ {
+                        tgid := findTgid[node.Children[i]]
+                        _, exists := visited[tgid]
+                        if !exists {
+                                visited[tgid] = true
+                                tgidNode, ok := treeMap.Load(tgid)
+                                if !ok {
+                                        continue
+                                }
+                                helloTree.Store(tgid, tgidNode)
+                                q.Enqueue(tgidNode)
+                        }
+                }
+        }
+
+        // TODO:
+        // 1.修改数据结构，使之自身即存储树结构，插入数据库后前端拿出来就能用
+        // 2.还有其余优化要做，比如线程退出时间与进程推出时间
+
+        // count := 0
+        // helloTree.Range(func(key, val interface{}) bool {
+        //      count++
+        //      fmt.Printf("tgid: %d\n", val.(treeNode).Tgid)
+        //      return true
+        // })
+        // fmt.Printf("Star: %d, res: %d\n", stared, count)
+
+        // 接下来处理文件
+}
+
+func ProMerge(a, b Process) (res Process) {
+        // 合并过程中会遇到什么问题？
+        res.Star = false
+
+        if a.StartTimestamp.IsZero() {
+                res.StartTimestamp = b.StartTimestamp
+        } else if b.StartTimestamp.IsZero() {
+                res.StartTimestamp = a.StartTimestamp
+        } else if a.StartTimestamp.Before(b.StartTimestamp) {
+                res.StartTimestamp = a.StartTimestamp
+        } else {
+                res.StartTimestamp = b.StartTimestamp
+        }
+
+        res.Ppid = a.Ppid
+        if a.ParentTgid == 0 {
+                res.ParentTgid = b.ParentTgid
+        } else {
+                res.ParentTgid = a.ParentTgid
+        }
+
+        res.Pid = a.Pid
+        if a.Tgid == 0 {
+                res.Tgid = b.Tgid
+        } else {
+                res.Tgid = a.Tgid
+        }
+
+        if len(a.Args) == 0 {
+                res.Args = b.Args
+        } else {
+                res.Args = a.Args
         }
 
-        for i := 0; i < len(res); i++ {
-                fmt.Printf("------\n%v\n", res[i])
+        if a.Comm == "" {
+                res.Comm = b.Comm
+        } else {
+                res.Comm = a.Comm
         }
+
+        if a.RootFS == "" {
+                res.RootFS = b.RootFS
+        } else {
+                res.RootFS = a.RootFS
+        }
+
+        if a.Cwd == "" {
+                res.Cwd = b.Cwd
+        } else {
+                res.Cwd = a.Cwd
+        }
+
+        res.Execve = append(a.Execve, b.Execve...)
+        res.Children = append(a.Children, b.Children...)
+
+        var flag bool // 真a假b
+        if a.ExitTimestamp.IsZero() {
+                flag = false
+        } else if b.ExitTimestamp.IsZero() {
+                flag = true
+        } else if a.ExitTimestamp.Before(b.ExitTimestamp) {
+                flag = true
+        } else {
+                flag = false
+        }
+
+        if flag {
+                res.ExitCode = a.ExitCode
+                res.ExitSignal = a.ExitSignal
+                res.ExitTimestamp = a.ExitTimestamp
+        } else {
+                res.ExitCode = b.ExitCode
+                res.ExitSignal = b.ExitSignal
+                res.ExitTimestamp = b.ExitTimestamp
+        }
+
+        return res
 }
diff --git a/filter/global.go b/filter/global.go
new file mode 100644
index 0000000..45706d4
--- /dev/null
+++ b/filter/global.go
@@ -0,0 +1,92 @@
+package main
+
+import (
+        "fmt"
+        "time"
+)
+
+type Exec struct {
+        Timestamp time.Time `bson:"timestamp"`
+        ExecArgs  []string  `bson:"execArgs"`
+}
+
+type Process struct {
+        Star           bool      `bson:"star"`
+        StartTimestamp time.Time `bson:"start_timestamp"`
+        Ppid           int       `bson:"ppid"`
+        ParentTgid     int       `bson:"parentTgid"`
+        Pid            int       `bson:"pid"`
+        Tgid           int       `bson:"tgid"`
+        Args           []string  `bson:"args"`
+        Comm           string    `bson:"comm"`
+        RootFS         string    `bson:"rootfs"`
+        Cwd            string    `bson:"cwd"`
+        Children       []int     `bson:"children"`
+        Execve         []Exec    `bson:"execve"`
+        ExitCode       int       `bson:"exit_code"`
+        ExitSignal     int       `bson:"exit_signal"`
+        ExitTimestamp  time.Time `bson:"exit_timestamp"`
+}
+
+func (p Process) String() string {
+        var res string
+        res = ""
+        res += fmt.Sprintf("timestamp\t%v\n", p.StartTimestamp)
+        res += fmt.Sprintf("ppid\t%d\nparentTgid\t%d\n", p.Ppid, p.ParentTgid)
+        res += fmt.Sprintf("pid\t%d\ntgid\t%d\nargs: ", p.Pid, p.Tgid)
+        for i := 0; i < len(p.Args); i++ {
+                res += fmt.Sprintf("%s ", p.Args[i])
+        }
+        res += fmt.Sprintf("\ncomm\t%s\ncwd\t%s\n", p.Comm, p.Cwd)
+        if len(p.Execve) != 0 {
+                res += fmt.Sprintf("exec:\n")
+                for i := 0; i < len(p.Execve); i++ {
+                        res += fmt.Sprintf("\ttimestamp: %v\n\texecArgs:\t", p.Execve[i].Timestamp)
+                        for j := 0; j < len(p.Execve[i].ExecArgs); j++ {
+                                res += fmt.Sprintf("%s ", p.Execve[i].ExecArgs[j])
+                        }
+                        res += fmt.Sprintf("\n")
+                }
+        }
+        res += fmt.Sprintf("children: ")
+        for i := 0; i < len(p.Children); i++ {
+                res += fmt.Sprintf("%d ", p.Children[i])
+        }
+        res += fmt.Sprintf("\n")
+        return res
+}
+
+// Queue 定义一个队列结构体
+type Queue struct {
+        items []interface{}
+}
+
+// NewQueue 创建一个新的队列
+func NewQueue() *Queue {
+        return &Queue{items: make([]interface{}, 0)}
+}
+
+// Enqueue 向队列中添加一个元素
+func (q *Queue) Enqueue(item interface{}) {
+        q.items = append(q.items, item)
+}
+
+// Dequeue 从队列中移除并返回队列前面的元素
+func (q *Queue) Dequeue() (interface{}, bool) {
+        if len(q.items) == 0 {
+                return nil, false
+        }
+        item := q.items[0]
+        q.items = q.items[1:]
+        return item, true
+}
+
+// Size 返回队列中的元素数量
+func (q *Queue) Size() int {
+        return len(q.items)
+}
+
+// IsEmpty 检查队列是否为空
+func (q *Queue) IsEmpty() bool {
+        return len(q.items) == 0
+}