Pārlūkot izejas kodu

add Shutdown to interfaces, rely tomb.Tomb for background routines killing

Thibault bui Koechlin 5 gadi atpakaļ
vecāks
revīzija
cedc67f56a
3 mainītis faili ar 26 papildinājumiem un 1 dzēšanām
  1. 10 0
      pkg/sqlite/commit.go
  2. 6 1
      pkg/sqlite/sqlite.go
  3. 10 0
      plugins/backend/sqlite.go

+ 10 - 0
pkg/sqlite/commit.go

@@ -35,6 +35,16 @@ func (c *Context) AutoCommit() {
 	ticker := time.NewTicker(200 * time.Millisecond)
 	for {
 		select {
+		case <-c.PusherTomb.Dying():
+			//we need to shutdown
+			log.Infof("sqlite routine shutdown")
+			if err := c.Flush(); err != nil {
+				log.Warningf("error while flushing records: %s", err)
+			}
+			if err := c.Db.Close(); err != nil {
+				log.Warningf("error while closing db : %s", err)
+			}
+			return
 		case <-ticker.C:
 			if atomic.LoadInt32(&c.count) != 0 &&
 				(atomic.LoadInt32(&c.count)%100 == 0 || time.Since(c.lastCommit) >= 500*time.Millisecond) {

+ 6 - 1
pkg/sqlite/sqlite.go

@@ -12,6 +12,7 @@ import (
 	"github.com/jinzhu/gorm"
 	_ "github.com/jinzhu/gorm/dialects/sqlite"
 	_ "github.com/mattn/go-sqlite3"
+	"gopkg.in/tomb.v2"
 )
 
 type Context struct {
@@ -21,6 +22,7 @@ type Context struct {
 	flush      bool
 	count      int32
 	lock       sync.Mutex //booboo
+	PusherTomb tomb.Tomb
 }
 
 func NewSQLite(cfg map[string]string) (*Context, error) {
@@ -62,6 +64,9 @@ func NewSQLite(cfg map[string]string) (*Context, error) {
 	if c.tx == nil {
 		return nil, fmt.Errorf("failed to begin sqlite transac : %s", err)
 	}
-	go c.AutoCommit()
+	c.PusherTomb.Go(func() error {
+		c.AutoCommit()
+		return nil
+	})
 	return c, nil
 }

+ 10 - 0
plugins/backend/sqlite.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"time"
 
 	"github.com/crowdsecurity/crowdsec/pkg/sqlite"
@@ -13,6 +14,15 @@ type pluginDB struct {
 	CTX *sqlite.Context
 }
 
+func (p *pluginDB) Shutdown() error {
+	p.CTX.PusherTomb.Kill(nil)
+	if err := p.CTX.PusherTomb.Wait(); err != nil {
+		return fmt.Errorf("DB shutdown error : %s", err)
+	}
+
+	return nil
+}
+
 func (p *pluginDB) Init(config map[string]string) error {
 	var err error
 	log.Debugf("sqlite config : %+v \n", config)