From ec3b4fc03a5bb1b27888b43c312756826f46db70 Mon Sep 17 00:00:00 2001
From: Bo-Yi Wu <appleboy.tw@gmail.com>
Date: Fri, 19 Aug 2016 21:16:10 +0800
Subject: [PATCH] Support create pid file.

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>
---
 config/config.go      | 29 ++++++++++++++++++++---------
 config/config.yml     |  4 ++++
 config/config_test.go |  8 ++++++++
 gorush.go             | 27 +++++++++++++++++++++++++++
 4 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/config/config.go b/config/config.go
index d5535bb..db1b883 100644
--- a/config/config.go
+++ b/config/config.go
@@ -18,15 +18,16 @@ type ConfYaml struct {
 
 // SectionCore is sub seciont of config.
 type SectionCore struct {
-	Port            string `yaml:"port"`
-	MaxNotification int    `yaml:"max_notification"`
-	WorkerNum       int    `yaml:"worker_num"`
-	QueueNum        int    `yaml:"queue_num"`
-	Mode            string `yaml:"mode"`
-	SSL             bool   `yaml:"ssl"`
-	CertPath        string `yaml:"cert_path"`
-	KeyPath         string `yaml:"key_path"`
-	HTTPProxy       string `yaml:"http_proxy"`
+	Port            string     `yaml:"port"`
+	MaxNotification int        `yaml:"max_notification"`
+	WorkerNum       int        `yaml:"worker_num"`
+	QueueNum        int        `yaml:"queue_num"`
+	Mode            string     `yaml:"mode"`
+	SSL             bool       `yaml:"ssl"`
+	CertPath        string     `yaml:"cert_path"`
+	KeyPath         string     `yaml:"key_path"`
+	HTTPProxy       string     `yaml:"http_proxy"`
+	PID             SectionPID `yaml:"pid"`
 }
 
 // SectionAPI is sub seciont of config.
@@ -88,6 +89,13 @@ type SectionBuntDB struct {
 	Path string `yaml:"path"`
 }
 
+// SectionPID is sub seciont of config.
+type SectionPID struct {
+	Enabled  bool   `yaml:"enabled"`
+	Path     string `yaml:"path"`
+	Override bool   `yaml:"override"`
+}
+
 // BuildDefaultPushConf is default config setting.
 func BuildDefaultPushConf() ConfYaml {
 	var conf ConfYaml
@@ -102,6 +110,9 @@ func BuildDefaultPushConf() ConfYaml {
 	conf.Core.KeyPath = "key.pem"
 	conf.Core.MaxNotification = 100
 	conf.Core.HTTPProxy = ""
+	conf.Core.PID.Enabled = false
+	conf.Core.PID.Path = "gorush.pid"
+	conf.Core.PID.Override = false
 
 	// Api
 	conf.API.PushURI = "/api/push"
diff --git a/config/config.yml b/config/config.yml
index 047c7a9..711b253 100644
--- a/config/config.yml
+++ b/config/config.yml
@@ -8,6 +8,10 @@ core:
   cert_path: "cert.pem"
   key_path: "key.pem"
   http_proxy: ""
+  pid:
+    enabled: false
+    path: "gorush.pid"
+    override: true
 
 api:
   push_uri: "/api/push"
diff --git a/config/config_test.go b/config/config_test.go
index 073c8da..fb59ddb 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -63,6 +63,10 @@ func (suite *ConfigTestSuite) TestValidateConfDefault() {
 	assert.Equal(suite.T(), "key.pem", suite.ConfGorushDefault.Core.KeyPath)
 	assert.Equal(suite.T(), 100, suite.ConfGorushDefault.Core.MaxNotification)
 	assert.Equal(suite.T(), "", suite.ConfGorushDefault.Core.HTTPProxy)
+	// Pid
+	assert.Equal(suite.T(), false, suite.ConfGorushDefault.Core.PID.Enabled)
+	assert.Equal(suite.T(), "gorush.pid", suite.ConfGorushDefault.Core.PID.Path)
+	assert.Equal(suite.T(), false, suite.ConfGorushDefault.Core.PID.Override)
 
 	// Api
 	assert.Equal(suite.T(), "/api/push", suite.ConfGorushDefault.API.PushURI)
@@ -111,6 +115,10 @@ func (suite *ConfigTestSuite) TestValidateConf() {
 	assert.Equal(suite.T(), "key.pem", suite.ConfGorush.Core.KeyPath)
 	assert.Equal(suite.T(), 100, suite.ConfGorush.Core.MaxNotification)
 	assert.Equal(suite.T(), "", suite.ConfGorush.Core.HTTPProxy)
+	// Pid
+	assert.Equal(suite.T(), false, suite.ConfGorush.Core.PID.Enabled)
+	assert.Equal(suite.T(), "gorush.pid", suite.ConfGorush.Core.PID.Path)
+	assert.Equal(suite.T(), true, suite.ConfGorush.Core.PID.Override)
 
 	// Api
 	assert.Equal(suite.T(), "/api/push", suite.ConfGorush.API.PushURI)
diff --git a/gorush.go b/gorush.go
index f593d7b..841c22b 100644
--- a/gorush.go
+++ b/gorush.go
@@ -7,6 +7,7 @@ import (
 	"github.com/appleboy/gorush/gorush"
 	"log"
 	"os"
+	"strconv"
 )
 
 func checkInput(token, message string) {
@@ -51,6 +52,28 @@ func usage() {
 	os.Exit(0)
 }
 
+func createPIDFile() error {
+	if !gorush.PushConf.Core.PID.Enabled {
+		return nil
+	}
+	_, err := os.Stat(gorush.PushConf.Core.PID.Path)
+	if os.IsNotExist(err) || gorush.PushConf.Core.PID.Override {
+		currentPid := os.Getpid()
+		file, err := os.Create(gorush.PushConf.Core.PID.Path)
+		if err != nil {
+			return fmt.Errorf("Can't create PID file: %v", err)
+		}
+		defer file.Close()
+		_, err = file.WriteString(strconv.FormatInt(int64(currentPid), 10))
+		if err != nil {
+			return fmt.Errorf("Can'write PID information on %s: %v", gorush.PushConf.Core.PID.Path, err)
+		}
+	} else {
+		return fmt.Errorf("%s already exists", gorush.PushConf.Core.PID.Path)
+	}
+	return nil
+}
+
 func main() {
 	opts := config.ConfYaml{}
 
@@ -207,6 +230,10 @@ func main() {
 		gorush.LogError.Fatal(err)
 	}
 
+	if err = createPIDFile(); err != nil {
+		gorush.LogError.Fatal(err)
+	}
+
 	gorush.InitAppStatus()
 	gorush.InitAPNSClient()
 	gorush.InitWorkers(gorush.PushConf.Core.WorkerNum, gorush.PushConf.Core.QueueNum)