@@ -18,12 +18,14 @@ import (
1818 "context"
1919 "flag"
2020 "fmt"
21+ "io"
2122 "log"
2223 "os"
2324 "path/filepath"
2425 "runtime/pprof"
2526 "time"
2627
28+ "github.com/felixge/fgprof"
2729 "github.com/sourcegraph/zoekt"
2830 "github.com/sourcegraph/zoekt/query"
2931 "github.com/sourcegraph/zoekt/shards"
@@ -76,11 +78,57 @@ func loadShard(fn string, verbose bool) (zoekt.Searcher, error) {
7678 return s , nil
7779}
7880
81+ func profile (path string , duration time.Duration , start func (io.Writer ) (stop func ())) func () bool {
82+ if path == "" {
83+ return func () bool { return false }
84+ }
85+
86+ f , err := os .Create (path )
87+ if err != nil {
88+ log .Fatal (err )
89+ }
90+
91+ t := time .Now ()
92+ stop := start (f )
93+
94+ return func () bool {
95+ if time .Since (t ) < duration {
96+ return true
97+ }
98+ stop ()
99+ f .Close ()
100+ return false
101+ }
102+ }
103+
104+ func startCPUProfile (path string , duration time.Duration ) func () bool {
105+ return profile (path , duration , func (w io.Writer ) func () {
106+ if err := pprof .StartCPUProfile (w ); err != nil {
107+ log .Fatal (err )
108+ }
109+
110+ return pprof .StopCPUProfile
111+ })
112+ }
113+
114+ func startFullProfile (path string , duration time.Duration ) func () bool {
115+ return profile (path , duration , func (w io.Writer ) func () {
116+ stop := fgprof .Start (w , fgprof .FormatPprof )
117+
118+ return func () {
119+ if err := stop (); err != nil {
120+ log .Fatal (err )
121+ }
122+ }
123+ })
124+ }
125+
79126func main () {
80127 shard := flag .String ("shard" , "" , "search in a specific shard" )
81128 index := flag .String ("index_dir" ,
82129 filepath .Join (os .Getenv ("HOME" ), ".zoekt" ), "search for index files in `directory`" )
83130 cpuProfile := flag .String ("cpu_profile" , "" , "write cpu profile to `file`" )
131+ fullProfile := flag .String ("full_profile" , "" , "write full profile to `file`" )
84132 profileTime := flag .Duration ("profile_time" , time .Second , "run this long to gather stats." )
85133 verbose := flag .Bool ("v" , false , "print some background data" )
86134 withRepo := flag .Bool ("r" , false , "print the repo before the file name" )
@@ -124,35 +172,19 @@ func main() {
124172
125173 var sOpts zoekt.SearchOptions
126174 sres , err := searcher .Search (context .Background (), query , & sOpts )
127- if * cpuProfile != "" {
128- // If profiling, do it another time so we measure with
129- // warm caches.
130- f , err := os .Create (* cpuProfile )
131- if err != nil {
132- log .Fatal (err )
133- }
134- defer f .Close ()
135- if * verbose {
136- log .Println ("Displaying matches..." )
137- }
138-
139- t := time .Now ()
140- if err := pprof .StartCPUProfile (f ); err != nil {
141- log .Fatal (err )
142- }
143- for {
144- sres , _ = searcher .Search (context .Background (), query , & sOpts )
145- if time .Since (t ) > * profileTime {
146- break
147- }
148- }
149- pprof .StopCPUProfile ()
150- }
151-
152175 if err != nil {
153176 log .Fatal (err )
154177 }
155178
179+ // If profiling, do it another time so we measure with
180+ // warm caches.
181+ for run := startCPUProfile (* cpuProfile , * profileTime ); run (); {
182+ sres , _ = searcher .Search (context .Background (), query , & sOpts )
183+ }
184+ for run := startFullProfile (* fullProfile , * profileTime ); run (); {
185+ sres , _ = searcher .Search (context .Background (), query , & sOpts )
186+ }
187+
156188 displayMatches (sres .Files , pat , * withRepo , * list )
157189 if * verbose {
158190 log .Printf ("stats: %#v" , sres .Stats )
0 commit comments