平滑退出

main.go

@@ -58,12 +58,15 @@ func main() {
 	case <-sigterm:
 		if braveConsumer != nil {
 			braveConsumer.Close()
+			glog.Info("braceConsumer平滑退出")
 		}
 		if healthConsumer != nil {
 			healthConsumer.Close()
+			glog.Info("healthConsumer平滑退出")
 		}
 		if db != nil {
 			db.Close()
+			glog.Info("dao平滑退出")
 		}
 		glog.Info(global.AppName + "平滑退出")
 	}

pkg/dao/dao.go

 package dao
 
 import (
+	"context"
 	"encoding/json"
 	"git.quantgroup.cn/DevOps/enoch/pkg/global"
 	"git.quantgroup.cn/DevOps/enoch/pkg/glog"
@@ -9,6 +10,7 @@ import (
 	"os"
 	"path"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -25,9 +27,14 @@ type Dao struct {
 	flashTime    time.Duration
 	dbAddress    string
 	cacheFileDir string
+	isClose      bool //平滑退出标记
+	ctx          context.Context
+	ctxCancel    context.CancelFunc
+	wg           *sync.WaitGroup
 }
 
 func New(batchSize int, flashTime time.Duration, dbAddress string, cacheFileDir string) *Dao {
+	ctx, cancel := context.WithCancel(context.Background())
 	rtn := &Dao{
 		batchSize:    batchSize,
 		size:         0,
@@ -35,6 +42,10 @@ func New(batchSize int, flashTime time.Duration, dbAddress string, cacheFileDir
 		flashTime:    flashTime,
 		dbAddress:    dbAddress,
 		cacheFileDir: cacheFileDir,
+		isClose:      false,
+		ctx:          ctx,
+		ctxCancel:    cancel,
+		wg:           new(sync.WaitGroup),
 	}
 
 	if stat, err := os.Stat(cacheFileDir); err != nil || !stat.IsDir() {
@@ -51,13 +62,18 @@ func New(batchSize int, flashTime time.Duration, dbAddress string, cacheFileDir
 }
 
 //平滑退出Dao，会flash缓存
-//TODO 未实现
 func (d *Dao) Close() {
-
+	d.isClose = true
+	d.ctxCancel()
+	d.wg.Wait()
 }
 
 func (d *Dao) flashFileCache() {
+	d.wg.Add(1)
+	timer := time.NewTimer(0)
 	for {
+		select {
+		case <-timer.C:
 			fileList, err := ioutil.ReadDir(d.cacheFileDir)
 			if err != nil || len(fileList) == 0 {
 				continue
@@ -102,7 +118,13 @@ func (d *Dao) flashFileCache() {
 					glog.Error("删除文件失败：", filePath)
 				}
 			}
-		time.Sleep(d.flashTime)
+			timer.Reset(d.flashTime)
+		case <-d.ctx.Done():
+			glog.Info("flash file cache 平滑退出")
+			timer.Stop()
+			d.wg.Done()
+			return
+		}
 	}
 }
 
@@ -111,11 +133,16 @@ func (d *Dao) MsgProcess(point *client.Point) {
 		return
 	}
 
+	if d.isClose {
+		return
+	}
+
 	d.channel <- point
 }
 
 //list满或者超时，则数据入库
 func (d *Dao) run() {
+	d.wg.Add(1)
 	defer func() {
 		if err := recover(); err != nil {
 			glog.Error(err)
@@ -138,6 +165,12 @@ func (d *Dao) run() {
 			go d.batchWrite(pointList)
 			pointList = make([]*client.Point, 0, d.batchSize)
 			timer.Reset(d.flashTime)
+		case <-d.ctx.Done():
+			go d.batchWrite(pointList)
+			glog.Info("存入influx平滑退出")
+			timer.Stop()
+			d.wg.Done()
+			return
 		}
 	}
 }
@@ -196,6 +229,9 @@ func (d *Dao) writeFile(pointList []*client.Point) error {
 }
 
 func (d *Dao) batchWrite(pointList []*client.Point) {
+	d.wg.Add(1)
+	defer d.wg.Done()
+
 	if len(pointList) == 0 {
 		return
 	}