github.com/grailbio/base@v0.0.11/grail/init.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache-2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package grail contains the Init function that all programs are expected to
     6  // call.
     7  package grail
     8  
     9  import (
    10  	"flag"
    11  	"os"
    12  	"strings"
    13  	"sync"
    14  
    15  	"github.com/google/gops/agent"
    16  	"github.com/grailbio/base/config"
    17  	"github.com/grailbio/base/log"
    18  	"github.com/grailbio/base/pprof"
    19  	"github.com/grailbio/base/shutdown"
    20  
    21  	// GRAIL applications require the AWS ticket provider.
    22  	_ "github.com/grailbio/base/config/awsticket"
    23  	"v.io/x/lib/vlog"
    24  )
    25  
    26  // Shutdown is a function that needs to be called to perform the final
    27  // cleanup.
    28  type Shutdown func()
    29  
    30  var (
    31  	initialized = false
    32  	mu          = sync.Mutex{}
    33  	gopsFlag    = flag.Bool("gops", false, "enable the gops listener")
    34  )
    35  
    36  // Init should be called once at the beginning at each executable that doesn't
    37  // use the github.com/grailbio/base/cmdutil. The Shutdown function should be called to
    38  // perform the final cleanup (closing logs for example).
    39  //
    40  // Init also applies a default configuration profile (see package
    41  // github.com/grailbio/base/config), and adds profile flags to the
    42  // default flag set. The default profile path used is $HOME/grail/profile.
    43  //
    44  // Note that this function will call flag.Parse().
    45  //
    46  // Suggested use:
    47  //
    48  //   shutdown := grail.Init()
    49  //   defer shutdown()
    50  func Init() Shutdown {
    51  	mu.Lock()
    52  	if initialized {
    53  		panic("Init called twice")
    54  	}
    55  	initialized = true
    56  	mu.Unlock()
    57  
    58  	profile := config.New()
    59  	config.NewDefault = func() *config.Profile {
    60  		if err := profile.Parse(strings.NewReader(defaultProfile)); err != nil {
    61  			panic("grail: error in default profile: " + err.Error())
    62  		}
    63  		if err := profile.ProcessFlags(); err != nil {
    64  			log.Fatal(err)
    65  		}
    66  		return profile
    67  	}
    68  	profile.RegisterFlags(flag.CommandLine, "", os.ExpandEnv("$HOME/grail/profile"))
    69  	flag.Parse()
    70  	if err := vlog.ConfigureLibraryLoggerFromFlags(); err != nil {
    71  		vlog.Error(err)
    72  	}
    73  	log.SetOutputter(VlogOutputter{})
    74  	if profile.NeedProcessFlags() {
    75  		_ = config.Application()
    76  	}
    77  
    78  	pprof.Start()
    79  	_, ok := os.LookupEnv("GOPS")
    80  	if ok || *gopsFlag {
    81  		if err := agent.Listen(agent.Options{}); err != nil {
    82  			log.Print(err)
    83  		}
    84  	}
    85  	return func() {
    86  		shutdown.Run()
    87  		pprof.Write(1)
    88  		vlog.FlushLog()
    89  	}
    90  }