From 96c60120132bdc67726bf486c039b73fda14f267 Mon Sep 17 00:00:00 2001
From: v2ray <admin@v2ray.com>
Date: Sun, 10 Jul 2016 23:11:42 +0200
Subject: [PATCH] refine main thread to capture all logs

---
 common/log/internal/log_writer.go | 17 ++++++++++++
 common/log/log.go                 |  5 ++++
 release/server/main.go            | 43 ++++++++++++++++++-------------
 3 files changed, 47 insertions(+), 18 deletions(-)

diff --git a/common/log/internal/log_writer.go b/common/log/internal/log_writer.go
index cc9fc64f..c3ad8a39 100644
--- a/common/log/internal/log_writer.go
+++ b/common/log/internal/log_writer.go
@@ -3,12 +3,15 @@ package internal
 import (
 	"log"
 	"os"
+	"time"
 
 	"github.com/v2ray/v2ray-core/common/platform"
+	"github.com/v2ray/v2ray-core/common/signal"
 )
 
 type LogWriter interface {
 	Log(LogEntry)
+	Close()
 }
 
 type NoOpLogWriter struct {
@@ -18,13 +21,18 @@ func (this *NoOpLogWriter) Log(entry LogEntry) {
 	entry.Release()
 }
 
+func (this *NoOpLogWriter) Close() {
+}
+
 type StdOutLogWriter struct {
 	logger *log.Logger
+	cancel *signal.CancelSignal
 }
 
 func NewStdOutLogWriter() LogWriter {
 	return &StdOutLogWriter{
 		logger: log.New(os.Stdout, "", log.Ldate|log.Ltime),
+		cancel: signal.NewCloseSignal(),
 	}
 }
 
@@ -33,10 +41,15 @@ func (this *StdOutLogWriter) Log(log LogEntry) {
 	log.Release()
 }
 
+func (this *StdOutLogWriter) Close() {
+	time.Sleep(500 * time.Millisecond)
+}
+
 type FileLogWriter struct {
 	queue  chan string
 	logger *log.Logger
 	file   *os.File
+	cancel *signal.CancelSignal
 }
 
 func (this *FileLogWriter) Log(log LogEntry) {
@@ -56,9 +69,12 @@ func (this *FileLogWriter) run() {
 		}
 		this.logger.Print(entry + platform.LineSeparator())
 	}
+	this.cancel.Done()
 }
 
 func (this *FileLogWriter) Close() {
+	close(this.queue)
+	<-this.cancel.WaitForDone()
 	this.file.Close()
 }
 
@@ -71,6 +87,7 @@ func NewFileLogWriter(path string) (*FileLogWriter, error) {
 		queue:  make(chan string, 16),
 		logger: log.New(file, "", log.Ldate|log.Ltime),
 		file:   file,
+		cancel: signal.NewCloseSignal(),
 	}
 	go logger.run()
 	return logger, nil
diff --git a/common/log/log.go b/common/log/log.go
index 82f99bd6..5fac4096 100644
--- a/common/log/log.go
+++ b/common/log/log.go
@@ -90,3 +90,8 @@ func Error(v ...interface{}) {
 		Values: v,
 	})
 }
+
+func Close() {
+	streamLoggerInstance.Close()
+	accessLoggerInstance.Close()
+}
diff --git a/release/server/main.go b/release/server/main.go
index 30072792..2e089c0b 100644
--- a/release/server/main.go
+++ b/release/server/main.go
@@ -43,15 +43,7 @@ func init() {
 	flag.StringVar(&configFile, "config", defaultConfigFile, "Config file for this Point server.")
 }
 
-func main() {
-	flag.Parse()
-
-	core.PrintVersion()
-
-	if *version {
-		return
-	}
-
+func startV2Ray() *point.Point {
 	switch *logLevel {
 	case "debug":
 		log.SetLogLevel(log.DebugLevel)
@@ -63,17 +55,17 @@ func main() {
 		log.SetLogLevel(log.ErrorLevel)
 	default:
 		fmt.Println("Unknown log level: " + *logLevel)
-		return
+		return nil
 	}
 
 	if len(configFile) == 0 {
 		log.Error("Config file is not set.")
-		return
+		return nil
 	}
 	config, err := point.LoadConfig(configFile)
 	if err != nil {
 		log.Error("Failed to read config file (", configFile, "): ", configFile, err)
-		return
+		return nil
 	}
 
 	if config.LogConfig != nil && len(config.LogConfig.AccessLog) > 0 {
@@ -83,23 +75,38 @@ func main() {
 	vPoint, err := point.NewPoint(config)
 	if err != nil {
 		log.Error("Failed to create Point server: ", err)
-		return
+		return nil
 	}
 
 	if *test {
 		fmt.Println("Configuration OK.")
-		return
+		return nil
 	}
 
 	err = vPoint.Start()
 	if err != nil {
 		log.Error("Error starting Point server: ", err)
+		return nil
+	}
+
+	return vPoint
+}
+
+func main() {
+	flag.Parse()
+
+	core.PrintVersion()
+
+	if *version {
 		return
 	}
 
-	osSignals := make(chan os.Signal, 1)
-	signal.Notify(osSignals, os.Interrupt, os.Kill)
+	if point := startV2Ray(); point != nil {
+		osSignals := make(chan os.Signal, 1)
+		signal.Notify(osSignals, os.Interrupt, os.Kill)
 
-	<-osSignals
-	vPoint.Close()
+		<-osSignals
+		point.Close()
+	}
+	log.Close()
 }