github.com/bhameyie/otto@v0.2.1-0.20160406174117-16052efa52ec/main.go (about) 1 package main 2 3 import ( 4 "fmt" 5 "io" 6 "io/ioutil" 7 "log" 8 "os" 9 "sync" 10 "time" 11 12 "github.com/hashicorp/otto/plugin" 13 "github.com/mitchellh/cli" 14 "github.com/mitchellh/panicwrap" 15 "github.com/mitchellh/prefixedio" 16 ) 17 18 func main() { 19 os.Exit(realMain()) 20 } 21 22 func realMain() int { 23 // Set a custom panicwrap cookie key and value. Since we're executing 24 // other panicwrap executables, the default cookie key/value cause 25 // weird errors if we don't change them. 26 var wrapConfig panicwrap.WrapConfig 27 wrapConfig.CookieKey = "OTTO_PANICWRAP_COOKIE" 28 wrapConfig.CookieValue = fmt.Sprintf( 29 "otto-%s-%s-%s", Version, VersionPrerelease, GitCommit) 30 31 if !panicwrap.Wrapped(&wrapConfig) { 32 // Determine where logs should go in general (requested by the user) 33 logWriter, err := logOutput() 34 if err != nil { 35 fmt.Fprintf(os.Stderr, "Couldn't setup log output: %s", err) 36 return 1 37 } 38 if logWriter == nil { 39 logWriter = ioutil.Discard 40 } 41 42 // We always send logs to a temporary file that we use in case 43 // there is a panic. Otherwise, we delete it. 44 logTempFile, err := ioutil.TempFile("", "otto-log") 45 if err != nil { 46 fmt.Fprintf(os.Stderr, "Couldn't setup logging tempfile: %s", err) 47 return 1 48 } 49 defer os.Remove(logTempFile.Name()) 50 defer logTempFile.Close() 51 52 // Tell the logger to log to this file 53 os.Setenv(EnvLog, "") 54 os.Setenv(EnvLogFile, "") 55 56 // Setup the prefixed readers that send data properly to 57 // stdout/stderr. 58 doneCh := make(chan struct{}) 59 outR, outW := io.Pipe() 60 go copyOutput(outR, doneCh) 61 62 // Create the configuration for panicwrap and wrap our executable 63 wrapConfig.Handler = panicHandler(logTempFile) 64 wrapConfig.Writer = io.MultiWriter(logTempFile, logWriter) 65 wrapConfig.Stdout = outW 66 exitStatus, err := panicwrap.Wrap(&wrapConfig) 67 if err != nil { 68 fmt.Fprintf(os.Stderr, "Couldn't start Otto: %s", err) 69 return 1 70 } 71 72 // If >= 0, we're the parent, so just exit 73 if exitStatus >= 0 { 74 // Close the stdout writer so that our copy process can finish 75 outW.Close() 76 77 // Wait for the output copying to finish 78 <-doneCh 79 80 return exitStatus 81 } 82 83 // We're the child, so just close the tempfile we made in order to 84 // save file handles since the tempfile is only used by the parent. 85 logTempFile.Close() 86 } 87 88 // Call the real main 89 return wrappedMain() 90 } 91 92 func wrappedMain() int { 93 // Make sure we cleanup any plugins that were launched. 94 defer plugin.CleanupClients() 95 96 log.SetOutput(os.Stderr) 97 log.Printf( 98 "[INFO] Otto version: %s %s %s", 99 Version, VersionPrerelease, GitCommit) 100 101 // Setup signal handlers 102 initSignalHandlers() 103 104 // Load the configuration 105 config := BuiltinConfig 106 107 // Run checkpoint 108 go runCheckpoint(&config) 109 110 // Get the command line args. We shortcut "--version" and "-v" to 111 // just show the version. 112 args := os.Args[1:] 113 for _, arg := range args { 114 if arg == "-v" || arg == "-version" || arg == "--version" { 115 newArgs := make([]string, len(args)+1) 116 newArgs[0] = "version" 117 copy(newArgs[1:], args) 118 args = newArgs 119 break 120 } 121 } 122 123 cli := &cli.CLI{ 124 Args: args, 125 Commands: Commands, 126 HelpFunc: cli.FilteredHelpFunc( 127 CommandsInclude, cli.BasicHelpFunc("otto")), 128 HelpWriter: os.Stdout, 129 } 130 131 exitCode, err := cli.Run() 132 if err != nil { 133 Ui.Error(fmt.Sprintf("Error executing CLI: %s", err.Error())) 134 return 1 135 } 136 137 return exitCode 138 } 139 140 // copyOutput uses output prefixes to determine whether data on stdout 141 // should go to stdout or stderr. This is due to panicwrap using stderr 142 // as the log and error channel. 143 func copyOutput(r io.Reader, doneCh chan<- struct{}) { 144 defer close(doneCh) 145 146 pr, err := prefixedio.NewReader(r) 147 if err != nil { 148 panic(err) 149 } 150 pr.FlushTimeout = 5 * time.Millisecond 151 152 stderrR, err := pr.Prefix(ErrorPrefix) 153 if err != nil { 154 panic(err) 155 } 156 stdoutR, err := pr.Prefix(OutputPrefix) 157 if err != nil { 158 panic(err) 159 } 160 defaultR, err := pr.Prefix("") 161 if err != nil { 162 panic(err) 163 } 164 165 var wg sync.WaitGroup 166 wg.Add(3) 167 go func() { 168 defer wg.Done() 169 io.Copy(os.Stderr, stderrR) 170 }() 171 go func() { 172 defer wg.Done() 173 io.Copy(os.Stdout, stdoutR) 174 }() 175 go func() { 176 defer wg.Done() 177 io.Copy(os.Stdout, defaultR) 178 }() 179 180 wg.Wait() 181 }