vendor/github.com/google/gopacket/examples/util/util.go

   1 // Copyright 2012 Google, Inc. All rights reserved.
   2 //
   3 // Use of this source code is governed by a BSD-style license
   4 // that can be found in the LICENSE file in the root of the source
   5 // tree.
   6
   7 // Package util provides shared utilities for all gopacket examples.
   8 package util
   9
  10 import (
  11         "flag"
  12         "log"
  13         "os"
  14         "runtime/pprof"
  15 )
  16
  17 var cpuprofile = flag.String("cpuprofile", "", "Where to write CPU profile")
  18
  19 // Run starts up stuff at the beginning of a main function, and returns a
  20 // function to defer until the function completes.  It should be used like this:
  21 //
  22 //   func main() {
  23 //     defer util.Run()()
  24 //     ... stuff ...
  25 //   }
  26 func Run() func() {
  27         flag.Parse()
  28         if *cpuprofile != "" {
  29                 f, err := os.Create(*cpuprofile)
  30                 if err != nil {
  31                         log.Fatalf("could not open cpu profile file %q", *cpuprofile)
  32                 }
  33                 pprof.StartCPUProfile(f)
  34                 return func() {
  35                         pprof.StopCPUProfile()
  36                         f.Close()
  37                 }
  38         }
  39         return func() {}
  40 }