From 257f4cb9050d29a38a992b814bd6e79e6f1bca99 Mon Sep 17 00:00:00 2001 From: Dmitry Vyukov Date: Wed, 16 Dec 2020 13:31:37 +0100 Subject: pkg/cmdprof: merge into pkg/tool cmdprof functionality seems to fit well into pkg/tool. --- pkg/cmdprof/cmdprof.go | 57 -------------------------------------------- pkg/tool/cmdprof.go | 44 ++++++++++++++++++++++++++++++++++ pkg/tool/flags.go | 20 +++++----------- pkg/tool/flags_test.go | 4 ++-- pkg/tool/tool.go | 15 ++++++++++++ sys/syz-sysgen/sysgen.go | 5 ++-- tools/syz-check/check.go | 5 ++-- tools/syz-cover/syz-cover.go | 5 ++-- 8 files changed, 73 insertions(+), 82 deletions(-) delete mode 100644 pkg/cmdprof/cmdprof.go create mode 100644 pkg/tool/cmdprof.go diff --git a/pkg/cmdprof/cmdprof.go b/pkg/cmdprof/cmdprof.go deleted file mode 100644 index ab0396d16..000000000 --- a/pkg/cmdprof/cmdprof.go +++ /dev/null @@ -1,57 +0,0 @@ -// Copyright 2020 syzkaller project authors. All rights reserved. -// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. - -// Package cmdprof simplifies cpu/memory profiling for command line tools. Use as: -// flag.Parse() -// defer cmdprof.Install()() -package cmdprof - -import ( - "flag" - "fmt" - "os" - "runtime" - "runtime/pprof" -) - -var ( - flagCPUProfile = flag.String("cpuprofile", "", "write CPU profile to this file") - flagMEMProfile = flag.String("memprofile", "", "write memory profile to this file") -) - -func Install() func() { - res := func() {} - failf := func(msg string, args ...interface{}) { - fmt.Fprintf(os.Stderr, msg+"\n", args...) - os.Exit(1) - } - if *flagCPUProfile != "" { - f, err := os.Create(*flagCPUProfile) - if err != nil { - failf("failed to create cpuprofile file: %v", err) - } - if err := pprof.StartCPUProfile(f); err != nil { - failf("failed to start cpu profile: %v", err) - } - res = func() { - pprof.StopCPUProfile() - f.Close() - } - } - if *flagMEMProfile != "" { - prev := res - res = func() { - prev() - f, err := os.Create(*flagMEMProfile) - if err != nil { - failf("failed to create memprofile file: %v", err) - } - defer f.Close() - runtime.GC() - if err := pprof.WriteHeapProfile(f); err != nil { - failf("failed to write mem profile: %v", err) - } - } - } - return res -} diff --git a/pkg/tool/cmdprof.go b/pkg/tool/cmdprof.go new file mode 100644 index 000000000..196a0cd24 --- /dev/null +++ b/pkg/tool/cmdprof.go @@ -0,0 +1,44 @@ +// Copyright 2020 syzkaller project authors. All rights reserved. +// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. + +package tool + +import ( + "os" + "runtime" + "runtime/pprof" +) + +// installProfiling simplifies cpu/memory profiling for command line tools. +func installProfiling(cpuprof, memprof string) func() { + res := func() {} + if cpuprof != "" { + f, err := os.Create(cpuprof) + if err != nil { + Failf("failed to create cpuprofile file: %v", err) + } + if err := pprof.StartCPUProfile(f); err != nil { + Failf("failed to start cpu profile: %v", err) + } + res = func() { + pprof.StopCPUProfile() + f.Close() + } + } + if memprof != "" { + prev := res + res = func() { + prev() + f, err := os.Create(memprof) + if err != nil { + Failf("failed to create memprofile file: %v", err) + } + defer f.Close() + runtime.GC() + if err := pprof.WriteHeapProfile(f); err != nil { + Failf("failed to write mem profile: %v", err) + } + } + } + return res +} diff --git a/pkg/tool/flags.go b/pkg/tool/flags.go index 197494823..0763cb502 100644 --- a/pkg/tool/flags.go +++ b/pkg/tool/flags.go @@ -8,39 +8,31 @@ import ( "encoding/hex" "flag" "fmt" - "os" "strings" "github.com/google/syzkaller/pkg/log" ) -// ParseFlags parses command line flags with flag.Parse and then applies optional flags created with OptionalFlags. -// This is intended for programmatic use only when we invoke older versions of binaries with new unsupported flags. -func ParseFlags() { - if err := parseFlags(flag.CommandLine, os.Args[1:]); err != nil { - Fail(err) - } -} - type Flag struct { Name string Value string } // OptionalFlags produces command line flag value that encapsulates the given flags as optional. -// Use ParseFlags to support optional flags in the binary. +// This is intended for programmatic use only when we invoke older versions of binaries with new unsupported flags. +// Use tool.Init to support optional flags in the binary. // The format keeps flags reasonably readable ("-optional=foo=bar:baz=123"), not subject to accidental splitting // into multiple arguments due to spaces and supports bool/non-bool flags. -func OptionalFlags(flags ...Flag) string { +func OptionalFlags(flags []Flag) string { return fmt.Sprintf("-%v=%v", optionalFlag, serializeFlags(flags)) } -func parseFlags(set *flag.FlagSet, args []string) error { - optional := set.String(optionalFlag, "", "optional flags for programmatic use only") +func ParseFlags(set *flag.FlagSet, args []string) error { + flagOptional := set.String(optionalFlag, "", "optional flags for programmatic use only") if err := set.Parse(args); err != nil { return err } - flags, err := deserializeFlags(*optional) + flags, err := deserializeFlags(*flagOptional) if err != nil { return err } diff --git a/pkg/tool/flags_test.go b/pkg/tool/flags_test.go index 6d52a4780..d818c6582 100644 --- a/pkg/tool/flags_test.go +++ b/pkg/tool/flags_test.go @@ -27,7 +27,7 @@ func TestParseFlags(t *testing.T) { {"", &Values{false, 1, "baz"}}, {"-foo -bar=2", &Values{true, 2, "baz"}}, {"-foo -bar=2 -qux", nil}, - {"-foo -bar=2 " + OptionalFlags(Flag{"qux", ""}), &Values{true, 2, "baz"}}, + {"-foo -bar=2 " + OptionalFlags([]Flag{{"qux", ""}}), &Values{true, 2, "baz"}}, } for i, test := range tests { t.Run(fmt.Sprint(i), func(t *testing.T) { @@ -41,7 +41,7 @@ func TestParseFlags(t *testing.T) { if args[0] == "" { args = args[1:] } - err := parseFlags(flags, args) + err := ParseFlags(flags, args) if test.vals == nil { if err == nil { t.Fatalf("parsing did not fail") diff --git a/pkg/tool/tool.go b/pkg/tool/tool.go index bb1f436cc..690bfa4b8 100644 --- a/pkg/tool/tool.go +++ b/pkg/tool/tool.go @@ -5,10 +5,25 @@ package tool import ( + "flag" "fmt" "os" ) +// Init handles common tasks for command line tools: +// - invokes flag.Parse +// - adds support for optional flags (see OptionalFlags) +// - adds support for cpu/mem profiling (-cpuprofile/memprofile flags) +// Use as defer tool.Init()(). +func Init() func() { + flagCPUProfile := flag.String("cpuprofile", "", "write CPU profile to this file") + flagMEMProfile := flag.String("memprofile", "", "write memory profile to this file") + if err := ParseFlags(flag.CommandLine, os.Args[1:]); err != nil { + Fail(err) + } + return installProfiling(*flagCPUProfile, *flagMEMProfile) +} + func Failf(msg string, args ...interface{}) { fmt.Fprintf(os.Stderr, msg+"\n", args...) os.Exit(1) diff --git a/sys/syz-sysgen/sysgen.go b/sys/syz-sysgen/sysgen.go index 62f36c3d3..b62e1dff2 100644 --- a/sys/syz-sysgen/sysgen.go +++ b/sys/syz-sysgen/sysgen.go @@ -18,11 +18,11 @@ import ( "text/template" "github.com/google/syzkaller/pkg/ast" - "github.com/google/syzkaller/pkg/cmdprof" "github.com/google/syzkaller/pkg/compiler" "github.com/google/syzkaller/pkg/hash" "github.com/google/syzkaller/pkg/osutil" "github.com/google/syzkaller/pkg/serializer" + "github.com/google/syzkaller/pkg/tool" "github.com/google/syzkaller/prog" "github.com/google/syzkaller/sys/targets" ) @@ -60,8 +60,7 @@ var srcDir = flag.String("src", "", "path to root of syzkaller source dir") var outDir = flag.String("out", "", "path to out dir") func main() { - flag.Parse() - defer cmdprof.Install()() + defer tool.Init()() var OSList []string for OS := range targets.List { diff --git a/tools/syz-check/check.go b/tools/syz-check/check.go index de7145dfb..5961cb6e8 100644 --- a/tools/syz-check/check.go +++ b/tools/syz-check/check.go @@ -37,10 +37,10 @@ import ( "unsafe" "github.com/google/syzkaller/pkg/ast" - "github.com/google/syzkaller/pkg/cmdprof" "github.com/google/syzkaller/pkg/compiler" "github.com/google/syzkaller/pkg/osutil" "github.com/google/syzkaller/pkg/symbolizer" + "github.com/google/syzkaller/pkg/tool" "github.com/google/syzkaller/prog" "github.com/google/syzkaller/sys/targets" ) @@ -59,8 +59,7 @@ func main() { fmt.Fprintf(os.Stderr, msg+"\n", args...) os.Exit(1) } - flag.Parse() - defer cmdprof.Install()() + defer tool.Init()() var warnings []Warn for arch, obj := range arches { if *obj == "" { diff --git a/tools/syz-cover/syz-cover.go b/tools/syz-cover/syz-cover.go index d019bfb78..2808cdd96 100644 --- a/tools/syz-cover/syz-cover.go +++ b/tools/syz-cover/syz-cover.go @@ -27,9 +27,9 @@ import ( "strconv" "strings" - "github.com/google/syzkaller/pkg/cmdprof" "github.com/google/syzkaller/pkg/cover" "github.com/google/syzkaller/pkg/osutil" + "github.com/google/syzkaller/pkg/tool" "github.com/google/syzkaller/sys/targets" ) @@ -43,8 +43,7 @@ func main() { flagKernelObj = flag.String("kernel_obj", "", "path to kernel build/obj dir") flagExport = flag.String("csv", "", "export coverage data in csv format (optional)") ) - flag.Parse() - defer cmdprof.Install()() + defer tool.Init()() if len(flag.Args()) == 0 { fmt.Fprintf(os.Stderr, "usage: syz-cover [flags] rawcover.file\n") -- cgit mrf-deployment