平滑退出

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,58 +62,69 @@ 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 {
-		fileList, err := ioutil.ReadDir(d.cacheFileDir)
-		if err != nil || len(fileList) == 0 {
-			continue
-		}
-		for _, file := range fileList {
-			sl := strings.Split(file.Name(), ":")
-			if len(sl) == 0 || sl[0] != filePrefixCache {
-				continue
-			}
-
-			//读取文件
-			filePath := path.Join(d.cacheFileDir, file.Name())
-			data, err := ioutil.ReadFile(filePath)
-			if err != nil {
-				glog.Error("can not read file:", filePath, err)
-				continue
-			}
-
-			cachePointList := make([]CachePoint, 0)
-			if err := json.Unmarshal(data, &cachePointList); err != nil {
-				glog.Error("can not unmarshal file:", filePath)
+		select {
+		case <-timer.C:
+			fileList, err := ioutil.ReadDir(d.cacheFileDir)
+			if err != nil || len(fileList) == 0 {
 				continue
 			}
+			for _, file := range fileList {
+				sl := strings.Split(file.Name(), ":")
+				if len(sl) == 0 || sl[0] != filePrefixCache {
+					continue
+				}
 
-			pointList := make([]*client.Point, 0)
-			for _, cachePoint := range cachePointList {
-				point, err := client.NewPoint(cachePoint.Name, cachePoint.Tags, cachePoint.Fields, cachePoint.Time)
+				//读取文件
+				filePath := path.Join(d.cacheFileDir, file.Name())
+				data, err := ioutil.ReadFile(filePath)
 				if err != nil {
+					glog.Error("can not read file:", filePath, err)
 					continue
 				}
-				pointList = append(pointList, point)
-			}
 
-			if err := d.writeDb(pointList); err != nil {
-				glog.Warn("flash file cache: can not write db")
-				continue
-			} else {
-				glog.Info("文件缓存写入db成功：", filePath)
-			}
+				cachePointList := make([]CachePoint, 0)
+				if err := json.Unmarshal(data, &cachePointList); err != nil {
+					glog.Error("can not unmarshal file:", filePath)
+					continue
+				}
+
+				pointList := make([]*client.Point, 0)
+				for _, cachePoint := range cachePointList {
+					point, err := client.NewPoint(cachePoint.Name, cachePoint.Tags, cachePoint.Fields, cachePoint.Time)
+					if err != nil {
+						continue
+					}
+					pointList = append(pointList, point)
+				}
+
+				if err := d.writeDb(pointList); err != nil {
+					glog.Warn("flash file cache: can not write db")
+					continue
+				} else {
+					glog.Info("文件缓存写入db成功：", filePath)
+				}
 
-			if err := os.Remove(filePath); err != nil {
-				glog.Error("删除文件失败：", filePath)
+				if err := os.Remove(filePath); err != nil {
+					glog.Error("删除文件失败：", filePath)
+				}
 			}
+			timer.Reset(d.flashTime)
+		case <-d.ctx.Done():
+			glog.Info("flash file cache 平滑退出")
+			timer.Stop()
+			d.wg.Done()
+			return
 		}
-		time.Sleep(d.flashTime)
 	}
 }
 
@@ -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
 	}