mirror of
https://github.com/sourcegraph/sourcegraph.git
synced 2026-02-06 17:31:43 +00:00
aa1121c6ba
This PR adds a new package memcmd, that adds a new abstraction called "Observer" that allows you to track the memory that a command (and all of its children) is using. (This package uses a polling approach with procfs, since [maxRSS on Linux is otherwise unreliable](https://jkz.wtf/random-linux-oddity-1-ru_maxrss) for our purposes). Example usage ```go import ( "context" "fmt" "os/exec" "time" "github.com/sourcegraph/sourcegraph/internal/memcmd" ) func Example() { const template = ` #!/usr/bin/env bash set -euo pipefail word=$(head -c "$((10 * 1024 * 1024))" </dev/zero | tr '\0' '\141') # 10MB worth of 'a's sleep 1 echo ${#word} ` cmd := exec.Command("bash", "-c", template) err := cmd.Start() if err != nil { panic(err) } observer, err := memcmd.NewLinuxObserver(context.Background(), cmd, 1*time.Millisecond) if err != nil { panic(err) } observer.Start() defer observer.Stop() err = cmd.Wait() if err != nil { panic(err) } memoryUsage, err := observer.MaxMemoryUsage() if err != nil { panic(err) } fmt.Println((0 < memoryUsage && memoryUsage < 50*1024*1024)) // Output should be between 0 and 50MB // Output: // true } ``` ## Test plan Unit tests Note that some tests only work on darwin, so you'll have to run those locally. ## Changelog This feature adds a package that allows us to track the memory usage of commands invoked via exec.Cmd. --------- Co-authored-by: Noah Santschi-Cooney <noah@santschi-cooney.ch>
68 lines
1.0 KiB
Go
68 lines
1.0 KiB
Go
//go:build linux
|
|
|
|
package memcmd_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"github.com/sourcegraph/sourcegraph/internal/memcmd"
|
|
)
|
|
|
|
func Example() {
|
|
const template = `
|
|
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
</dev/zero head -c $((1024**2*50)) | tail
|
|
sleep 1
|
|
`
|
|
|
|
tempDir, err := os.MkdirTemp("", "foo")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
defer func() {
|
|
_ = os.RemoveAll(tempDir)
|
|
}()
|
|
|
|
p := filepath.Join(tempDir, "/script.sh")
|
|
err = os.WriteFile(p, []byte(template), 0755)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
cmd := exec.Command("bash", "-c", p) // 50MB
|
|
err = cmd.Start()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
observer, err := memcmd.NewLinuxObserver(context.Background(), cmd, 1*time.Millisecond)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
observer.Start()
|
|
defer observer.Stop()
|
|
|
|
err = cmd.Wait()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
memoryUsage, err := observer.MaxMemoryUsage()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
fmt.Println((0 < memoryUsage && memoryUsage < 100*1024*1024)) // Output should be between 0 and 100MB
|
|
|
|
// Output:
|
|
// true
|
|
}
|