From 2abb5c96dc7ebe4bae75e8ddc6611eb1fa173495 Mon Sep 17 00:00:00 2001 From: mozillazg Date: Sat, 24 Aug 2024 06:58:51 +0000 Subject: [PATCH 1/2] chore(user): Optimize initialization speed --- internal/metadata/net.go | 2 +- internal/metadata/process.go | 66 ++++++++++++++++++++++++------------ 2 files changed, 46 insertions(+), 22 deletions(-) diff --git a/internal/metadata/net.go b/internal/metadata/net.go index e73358fa..9ef3c54b 100644 --- a/internal/metadata/net.go +++ b/internal/metadata/net.go @@ -37,7 +37,7 @@ func GetCurrentConnects(ctx context.Context, pids []int, all bool) ([]Connection } } for _, stat := range stats { - if stat.Laddr.Port == 0 || stat.Raddr.Port == 0 || stat.Status != "ESTABLISHED" { + if stat.Pid == 0 || stat.Laddr.Port == 0 || stat.Raddr.Port == 0 || stat.Status != "ESTABLISHED" { continue } conn, err := convertConnectionStat(stat) diff --git a/internal/metadata/process.go b/internal/metadata/process.go index 208a1d85..edb03349 100644 --- a/internal/metadata/process.go +++ b/internal/metadata/process.go @@ -3,6 +3,8 @@ package metadata import ( "context" "fmt" + "runtime" + "sort" "sync" "time" @@ -49,6 +51,7 @@ func (c *ProcessCache) WithContainerCache(cc *ContainerCache) *ProcessCache { func (c *ProcessCache) Start(ctx context.Context) { // TODO: change to get running processes via ebpf task iter + log.Info("start to fill running process info") if err := c.fillRunningProcesses(ctx); err != nil { log.Errorf("fill running processes info failed: %s", err) } @@ -56,36 +59,54 @@ func (c *ProcessCache) Start(ctx context.Context) { } func (c *ProcessCache) fillRunningProcesses(ctx context.Context) error { + log.Info("start to get all processes") ps, err := process.ProcessesWithContext(ctx) if err != nil { return fmt.Errorf(": %w", err) } + sort.Slice(ps, func(i, j int) bool { + return ps[i].Pid < ps[j].Pid + }) + + log.Info("start to add process events with these processes data") + pool := make(chan struct{}, runtime.NumCPU()) + wg := sync.WaitGroup{} for _, p := range ps { + p := p if p.Pid == 0 { continue } - ppid := 0 - if parent, err := p.ParentWithContext(ctx); err == nil { - ppid = int(parent.Pid) - } - filename, _ := p.Exe() - if filename == "" { - filename, _ = p.Name() - } - args, _ := p.CmdlineSlice() - e := event.ProcessExec{ - PPid: ppid, - Pid: int(p.Pid), - Filename: filename, - FilenameTruncated: false, - Args: args, - ArgsTruncated: false, - PidNs: utils.GetPidNamespaceFromPid(int(p.Pid)), - MntNs: utils.GetMountNamespaceFromPid(int(p.Pid)), - Netns: utils.GetNetworkNamespaceFromPid(int(p.Pid)), - } - c.AddItem(e) + pool <- struct{}{} + wg.Add(1) + go func(p *process.Process) { + defer func() { + <-pool + wg.Done() + }() + ppid := 0 + if parent, err := p.ParentWithContext(ctx); err == nil { + ppid = int(parent.Pid) + } + filename, _ := p.Exe() + if filename == "" { + filename, _ = p.Name() + } + args, _ := p.CmdlineSlice() + e := event.ProcessExec{ + PPid: ppid, + Pid: int(p.Pid), + Filename: filename, + FilenameTruncated: false, + Args: args, + ArgsTruncated: false, + PidNs: utils.GetPidNamespaceFromPid(int(p.Pid)), + MntNs: utils.GetMountNamespaceFromPid(int(p.Pid)), + Netns: utils.GetNetworkNamespaceFromPid(int(p.Pid)), + } + c.AddItem(e) + }(p) } + wg.Wait() return nil } @@ -140,6 +161,9 @@ func (c *ProcessCache) jobCleanDead() { func (c *ProcessCache) AddItemWithContext(exec event.ProcessExec, rawCtx types.PacketContext) { pid := exec.Pid + if pid == 0 { + return + } // if exec.CgroupName == "" || // strings.HasSuffix(exec.CgroupName, ".slice") || From 44a144325cdfb948d0f6ac1c9e0172f19e9f5fd6 Mon Sep 17 00:00:00 2001 From: mozillazg Date: Sat, 24 Aug 2024 07:16:18 +0000 Subject: [PATCH 2/2] feat(user): Add `--micro`, `--nano`, and `--time-stamp-precision` flags for time stamp precision control --- cmd/options.go | 22 ++++++++++++++++------ cmd/read.go | 1 + cmd/root.go | 6 ++++++ cmd/writer.go | 1 + internal/writer/stdout.go | 23 ++++++++++++++--------- 5 files changed, 38 insertions(+), 15 deletions(-) diff --git a/cmd/options.go b/cmd/options.go index c8a03d7f..f33bd487 100644 --- a/cmd/options.go +++ b/cmd/options.go @@ -31,12 +31,18 @@ type Options struct { printPacketNumber bool dontPrintTimestamp bool onlyPrintCount bool - dontConvertAddr int - verbose int - containerId string - containerName string - podName string - podNamespace string + + timeStampPrecision string + timeStampMicro bool + timeStampNano bool + + dontConvertAddr int + verbose int + + containerId string + containerName string + podName string + podNamespace string eventChanSize uint delayBeforeHandlePacketEvents time.Duration @@ -79,6 +85,10 @@ func (o Options) DirectionInOut() bool { return o.direction == "inout" } +func (o Options) TimeStampAsNano() bool { + return o.timeStampNano || o.timeStampPrecision == "nano" +} + func prepareOptions(opts *Options, rawArgs []string, args []string) { subProgArgs := getSubProgArgs(rawArgs) opts.pcapFilter = strings.Join(args, " ") diff --git a/cmd/read.go b/cmd/read.go index 1db72cff..2b7f76ee 100644 --- a/cmd/read.go +++ b/cmd/read.go @@ -26,6 +26,7 @@ func read(ctx context.Context, opts Options) error { stdoutWriter.OneLine = opts.oneLine stdoutWriter.PrintNumber = opts.printPacketNumber stdoutWriter.NoTimestamp = opts.dontPrintTimestamp + stdoutWriter.TimestampNano = opts.TimeStampAsNano() if opts.onlyPrintCount { stdoutWriter.DoNothing = true } diff --git a/cmd/root.go b/cmd/root.go index 5873a432..981a8e2a 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -105,6 +105,12 @@ func init() { "/var/lib/ptcpdump/btf/$(uname -r).btf", "download BTF file from https://mirrors.openanolis.cn/coolbpf/btf/ and https://github.com/aquasecurity/btfhub-archive/"}, ", ")), ) + rootCmd.Flags().StringVar(&opts.timeStampPrecision, "time-stamp-precision", "micro", + "When capturing, set the time stamp precision for the capture to the format") + rootCmd.Flags().BoolVar(&opts.timeStampMicro, "micro", false, + "Shorthands for --time-stamp-precision=micro") + rootCmd.Flags().BoolVar(&opts.timeStampNano, "nano", false, + "Shorthands for --time-stamp-precision=nano") silenceKlog() } diff --git a/cmd/writer.go b/cmd/writer.go index 1e40987f..1f898c73 100644 --- a/cmd/writer.go +++ b/cmd/writer.go @@ -61,6 +61,7 @@ func getWriters(opts Options, pcache *metadata.ProcessCache) ([]writer.PacketWri stdoutWriter.OneLine = opts.oneLine stdoutWriter.PrintNumber = opts.printPacketNumber stdoutWriter.NoTimestamp = opts.dontPrintTimestamp + stdoutWriter.TimestampNano = opts.TimeStampAsNano() if opts.verbose >= 1 { stdoutWriter.FormatStyle = pktdump.FormatStyleVerbose } diff --git a/internal/writer/stdout.go b/internal/writer/stdout.go index c92f61d8..5b9a1ced 100644 --- a/internal/writer/stdout.go +++ b/internal/writer/stdout.go @@ -14,14 +14,15 @@ import ( ) type StdoutWriter struct { - pcache *metadata.ProcessCache - w io.Writer - Decoder gopacket.Decoder - OneLine bool - PrintNumber bool - NoTimestamp bool - DoNothing bool - FormatStyle pktdump.FormatStyle + pcache *metadata.ProcessCache + w io.Writer + Decoder gopacket.Decoder + OneLine bool + PrintNumber bool + NoTimestamp bool + TimestampNano bool + DoNothing bool + FormatStyle pktdump.FormatStyle n int64 } @@ -88,7 +89,11 @@ func (w *StdoutWriter) Write(e *event.Packet) error { } if !w.NoTimestamp { - builder.WriteString(fmt.Sprintf("%s ", e.Time.Local().Format("15:04:05.000000"))) + layout := "15:04:05.000000" + if w.TimestampNano { + layout = "15:04:05.000000000" + } + builder.WriteString(fmt.Sprintf("%s ", e.Time.Local().Format(layout))) } if ifName != "" {