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

This commit is contained in:
Thibault bui Koechlin 2020-06-17 01:11:59 +02:00
parent e9c16ee064
commit cedc67f56a
3 changed files with 26 additions and 1 deletions

View file

@ -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) {

View file

@ -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
}

View file

@ -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)