demonstrating polling technique for finding memory usage of child process on linux

main.go

package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/exec"
	"sync"
	"syscall"
	"time"

	"github.com/hashicorp/go-multierror"
	"github.com/inhies/go-bytesize"
	"github.com/prometheus/procfs"
)

const program = `
#!/bin/bash
python -c 'import time; x = bytearray(1024 * 1024 * 1024); print("allocated this many bytes: " + str(len(x)))'
echo "done" # force bash to fork
`

func main() {
	ctx := context.Background()

	cmd := exec.CommandContext(ctx, "bash", "-c", program)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setpgid: true,
	}

	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	memoryUsage, err := startWithMemoryObserver(cmd, 100*time.Millisecond)
	if err != nil {
		log.Fatalf("failed to run command: %v", err)
	}

	log.Printf("max memory usage: %d (%s)", memoryUsage, bytesize.ByteSize(memoryUsage))

}

func reportCumulativeMemoryUsageForPIDAndChildren(ctx context.Context, pid int) (memoryUsedBytes uint64, err error) {
	allRSS := uint64(0)
	var errs *multierror.Error

	procStack := []int{pid}
	for len(procStack) > 0 {
		select {
		case <-ctx.Done():
			return allRSS, ctx.Err()
		default:
		}

		current := procStack[0]
		procStack = procStack[1:]

		rss, err := reportMemoryUsage(current)
		if err != nil {
			if !os.IsNotExist(err) {
				err = fmt.Errorf("failed to report memory usage for pid %d: %w", current, err)
				errs = multierror.Append(errs, err)
			}

			continue
		}

		allRSS += rss

		procs, err := procfs.AllProcs()
		if err != nil {
			if !os.IsNotExist(err) {
				err = fmt.Errorf("failed to list all processes: %w", err)
				errs = multierror.Append(errs, err)
			}

			continue
		}

		for _, p := range procs {
			stat, err := p.Stat()
			if err != nil {
				if !os.IsNotExist(err) { // Ignore non-longer-existent processes
					err = fmt.Errorf("failed to get stats for pid %d: %w", p.PID, err)
					errs = multierror.Append(errs, err)
				}

				continue
			}

			if stat.PPID == current {
				// This a child of the original process, so
				// we need to report its memory usage as well.
				procStack = append(procStack, p.PID)
			}
		}
	}

	return allRSS, errs.ErrorOrNil()
}

func reportMemoryUsage(pid int) (rss uint64, err error) {
	fs, err := procfs.NewProc(pid)
	if err != nil {
		return 0, fmt.Errorf("failed to get procfs: %w", err)
	}

	status, err := fs.NewStatus()
	if err != nil {
		return 0, fmt.Errorf("failed to get status: %w", err)
	}

	return status.VmRSS, nil
}

func startWithMemoryObserver(cmd *exec.Cmd, sampleInterval time.Duration) (maxMemoryBytes uint64, err error) {
	err = cmd.Start()
	if err != nil {
		return 0, fmt.Errorf("failed to start command: %w", err)
	}

	pid := cmd.Process.Pid
	log.Printf("started command with pid %d", pid)

	maxMemoryUsage := uint64(0)

	var wg sync.WaitGroup
	wg.Add(1)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go func() {
		ticker := time.NewTicker(sampleInterval)
		defer func() {
			ticker.Stop()
			wg.Done()
		}()

		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				memoryUsage, err := reportCumulativeMemoryUsageForPIDAndChildren(ctx, pid)
				if err != nil {
					log.Printf("encountered error while reporting memory usage: %v", err) // Ignore
				}

				if memoryUsage > maxMemoryUsage {
					maxMemoryUsage = memoryUsage
				}
			}
		}
	}()

	return maxMemoryUsage, cmd.Wait()
}

https://jkz.wtf/random-linux-oddity-1-ru_maxrss

2024/05/14 16:56:28 started command with pid 14934
allocated this many bytes: 1073741824
2024/05/14 16:56:28 max memory usage: 995098624 (949.00MB)

Process finished with the exit code 0