diff --git a/pkg/sqlite/commit.go b/pkg/sqlite/commit.go index 4c51c6780..657db560b 100644 --- a/pkg/sqlite/commit.go +++ b/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) { diff --git a/pkg/sqlite/sqlite.go b/pkg/sqlite/sqlite.go index 9234dc66d..b3c5d8a63 100644 --- a/pkg/sqlite/sqlite.go +++ b/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 } diff --git a/plugins/backend/sqlite.go b/plugins/backend/sqlite.go index ec6dfe4be..4e1943018 100644 --- a/plugins/backend/sqlite.go +++ b/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)