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) ticker := time.NewTicker(200 * time.Millisecond)
for { for {
select { 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: case <-ticker.C:
if atomic.LoadInt32(&c.count) != 0 && if atomic.LoadInt32(&c.count) != 0 &&
(atomic.LoadInt32(&c.count)%100 == 0 || time.Since(c.lastCommit) >= 500*time.Millisecond) { (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"
_ "github.com/jinzhu/gorm/dialects/sqlite" _ "github.com/jinzhu/gorm/dialects/sqlite"
_ "github.com/mattn/go-sqlite3" _ "github.com/mattn/go-sqlite3"
"gopkg.in/tomb.v2"
) )
type Context struct { type Context struct {
@ -21,6 +22,7 @@ type Context struct {
flush bool flush bool
count int32 count int32
lock sync.Mutex //booboo lock sync.Mutex //booboo
PusherTomb tomb.Tomb
} }
func NewSQLite(cfg map[string]string) (*Context, error) { func NewSQLite(cfg map[string]string) (*Context, error) {
@ -62,6 +64,9 @@ func NewSQLite(cfg map[string]string) (*Context, error) {
if c.tx == nil { if c.tx == nil {
return nil, fmt.Errorf("failed to begin sqlite transac : %s", err) 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 return c, nil
} }

View file

@ -1,6 +1,7 @@
package main package main
import ( import (
"fmt"
"time" "time"
"github.com/crowdsecurity/crowdsec/pkg/sqlite" "github.com/crowdsecurity/crowdsec/pkg/sqlite"
@ -13,6 +14,15 @@ type pluginDB struct {
CTX *sqlite.Context 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 { func (p *pluginDB) Init(config map[string]string) error {
var err error var err error
log.Debugf("sqlite config : %+v \n", config) log.Debugf("sqlite config : %+v \n", config)