vitess.io/vitess@v0.16.2/go/exit/exit.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  /*
    18  Package exit provides an alternative to os.Exit(int).
    19  
    20  Unlike os.Exit(int), exit.Return(int) will run deferred functions before
    21  terminating. It's effectively like a return from main(), except you can specify
    22  the exit code.
    23  
    24  Defer a call to exit.Recover() or exit.RecoverAll() at the beginning of main().
    25  Use exit.Return(int) to initiate an exit.
    26  
    27  	func main() {
    28  		defer exit.Recover()
    29  		defer cleanup()
    30  		...
    31  		if err != nil {
    32  			// Return from main() with a non-zero exit code,
    33  			// making sure to run deferred cleanup.
    34  			exit.Return(1)
    35  		}
    36  		...
    37  	}
    38  
    39  All functions deferred *after* defer exit.Recover()/RecoverAll() will be
    40  executed before the exit. This is why the defer for this package should
    41  be the first statement in main().
    42  
    43  NOTE: This mechanism only works if exit.Return() is called from the same
    44  goroutine that deferred exit.Recover(). Usually this means Return() should
    45  only be used from within main(), or within functions that are only ever
    46  called from main(). See Recover() and Return() for more details.
    47  */
    48  package exit
    49  
    50  import (
    51  	"os"
    52  
    53  	"vitess.io/vitess/go/tb"
    54  	"vitess.io/vitess/go/vt/log"
    55  )
    56  
    57  type exitCode int
    58  
    59  var (
    60  	exitFunc = os.Exit // can be faked out for testing
    61  )
    62  
    63  // Recover should be deferred as the first line of main(). It recovers the
    64  // panic initiated by Return and converts it to a call to os.Exit. Any
    65  // functions deferred after Recover in the main goroutine will be executed
    66  // prior to exiting. Recover will re-panic anything other than the panic it
    67  // expects from Return.
    68  func Recover() {
    69  	doRecover(recover(), false)
    70  }
    71  
    72  // RecoverAll can be deferred instead of Recover as the first line of main().
    73  // Instead of re-panicking, RecoverAll will absorb any panic and convert it to
    74  // an error log entry with a stack trace, followed by a call to os.Exit(255).
    75  func RecoverAll() {
    76  	doRecover(recover(), true)
    77  }
    78  
    79  func doRecover(err any, recoverAll bool) {
    80  	if err == nil {
    81  		return
    82  	}
    83  
    84  	switch code := err.(type) {
    85  	case exitCode:
    86  		exitFunc(int(code))
    87  	default:
    88  		if recoverAll {
    89  			log.Errorf("panic: %v", tb.Errorf("%v", err))
    90  			exitFunc(255)
    91  		} else {
    92  			panic(err)
    93  		}
    94  	}
    95  }
    96  
    97  // Return initiates a panic that sends the return code to the deferred Recover,
    98  // executing other deferred functions along the way. When the panic reaches
    99  // Recover, the return code will be passed to os.Exit. This should only be
   100  // called from the main goroutine.
   101  func Return(code int) {
   102  	panic(exitCode(code))
   103  }