From 0439cdfab5d02c3ab9d3c13dc254ead731bc602b Mon Sep 17 00:00:00 2001
From: Josh Rickmar <jrick@conformal.com>
Date: Sat, 21 Jun 2014 10:39:27 -0500
Subject: [PATCH] Flush logs and run other defers before os.Exit.

As calls os.Exit do not run deferred functions (such as log flushing),
the real main function should simply run a main helper function that,
rather than exiting the program, runs all defers and returns a
possibly non-nil error.  The real main function can then check the
error and close the program with an error exit status when a fatal
error occured.
---
 cmd.go     | 67 +++++++++++++++++++++++++++++++-----------------------
 updates.go | 19 ++++++++--------
 2 files changed, 48 insertions(+), 38 deletions(-)

diff --git a/cmd.go b/cmd.go
index 763d0dd..53f796d 100644
--- a/cmd.go
+++ b/cmd.go
@@ -102,37 +102,24 @@ func accessClient() (*rpcClient, error) {
 	return c, nil
 }
 
-func clientConnect(certs []byte, newClient chan<- *rpcClient) {
-	const initialWait = 5 * time.Second
-	wait := initialWait
-	for {
-		client, err := newRPCClient(certs)
-		if err != nil {
-			log.Warnf("Unable to open chain server client "+
-				"connection: %v", err)
-			time.Sleep(wait)
-			wait <<= 1
-			if wait > time.Minute {
-				wait = time.Minute
-			}
-			continue
-		}
-
-		wait = initialWait
-
-		client.Start()
-		newClient <- client
-
-		client.WaitForShutdown()
+func main() {
+	// Work around defer not working after os.Exit.
+	if err := walletMain(); err != nil {
+		os.Exit(1)
 	}
 }
 
-func main() {
+// walletMain is a work-around main function that is required since deferred
+// functions (such as log flushing) are not called with calls to os.Exit.
+// Instead, main runs this function and checks for a non-nil error, at which
+// point any defers have already run, and if the error is non-nil, the program
+// can be exited with an error exit status.
+func walletMain() error {
 	// Load configuration and parse command line.  This function also
 	// initializes logging and configures it accordingly.
 	tcfg, _, err := loadConfig()
 	if err != nil {
-		os.Exit(1)
+		return err
 	}
 	cfg = tcfg
 	defer backendLog.Flush()
@@ -152,11 +139,14 @@ func main() {
 	certs, err := ioutil.ReadFile(cfg.CAFile)
 	if err != nil {
 		log.Errorf("cannot open CA file: %v", err)
-		os.Exit(1)
+		return err
 	}
 
 	// Check and update any old file locations.
-	updateOldFileLocations()
+	err = updateOldFileLocations()
+	if err != nil {
+		return err
+	}
 
 	// Start account manager and open accounts.
 	AcctMgr.Start()
@@ -164,7 +154,7 @@ func main() {
 	server, err = newRPCServer(cfg.SvrListeners)
 	if err != nil {
 		log.Errorf("Unable to create HTTP server: %v", err)
-		os.Exit(1)
+		return err
 	}
 
 	// Start HTTP server to serve wallet client connections.
@@ -174,5 +164,26 @@ func main() {
 	// reconnections if the client could not be successfully connected.
 	clientChan := make(chan *rpcClient)
 	go clientAccess(clientChan)
-	clientConnect(certs, clientChan)
+	const initialWait = 5 * time.Second
+	wait := initialWait
+	for {
+		client, err := newRPCClient(certs)
+		if err != nil {
+			log.Warnf("Unable to open chain server client "+
+				"connection: %v", err)
+			time.Sleep(wait)
+			wait <<= 1
+			if wait > time.Minute {
+				wait = time.Minute
+			}
+			continue
+		}
+
+		wait = initialWait
+
+		client.Start()
+		clientChan <- client
+
+		client.WaitForShutdown()
+	}
 }
diff --git a/updates.go b/updates.go
index 9db6a24..87dc984 100644
--- a/updates.go
+++ b/updates.go
@@ -29,10 +29,7 @@ var ErrNotAccountDir = errors.New("directory is not an account directory")
 
 // updateOldFileLocations moves files for wallets, transactions, and
 // recorded unspent transaction outputs to more recent locations.
-//
-// If any errors are encounted during this function, the application is
-// closed.
-func updateOldFileLocations() {
+func updateOldFileLocations() error {
 	// Before version 0.1.1, accounts were saved with the following
 	// format:
 	//
@@ -69,7 +66,7 @@ func updateOldFileLocations() {
 
 	datafi, err := os.Open(cfg.DataDir)
 	if err != nil {
-		return
+		return nil
 	}
 	defer func() {
 		if err := datafi.Close(); err != nil {
@@ -81,7 +78,7 @@ func updateOldFileLocations() {
 	fi, err := datafi.Readdir(0)
 	if err != nil {
 		log.Errorf("Cannot read files in data directory: %v", err)
-		os.Exit(1)
+		return err
 	}
 
 	acctsExist := false
@@ -97,14 +94,14 @@ func updateOldFileLocations() {
 		}
 	}
 	if !acctsExist {
-		return
+		return nil
 	}
 
 	// Create testnet directory, if it doesn't already exist.
 	netdir := filepath.Join(cfg.DataDir, "testnet")
 	if err := checkCreateDir(netdir); err != nil {
 		log.Errorf("Cannot continue without a testnet directory: %v", err)
-		os.Exit(1)
+		return err
 	}
 
 	// Check all files in the datadir for old accounts to update.
@@ -124,7 +121,7 @@ func updateOldFileLocations() {
 
 		default: // all other non-nil errors
 			log.Errorf("Cannot open old account directory: %v", err)
-			os.Exit(1)
+			return err
 		}
 
 		log.Infof("Updating old file locations for account %v", account)
@@ -136,7 +133,7 @@ func updateOldFileLocations() {
 			if err := Rename(old, new); err != nil {
 				log.Errorf("Cannot move old %v for account %v to new location: %v",
 					"wallet.bin", account, err)
-				os.Exit(1)
+				return err
 			}
 		}
 
@@ -145,6 +142,8 @@ func updateOldFileLocations() {
 			log.Warnf("Could not remove pre 0.1.1 account directory: %v", err)
 		}
 	}
+
+	return nil
 }
 
 type oldAccountDir struct {