1
0
mirror of https://github.com/netsampler/goflow2.git synced 2024-05-06 15:54:52 +00:00

feature: implement SIGHUP for log rotation when using transport.destination

This commit is contained in:
lspgn
2021-05-27 19:33:19 -07:00
parent df01c790df
commit 686497ea5c

View File

@ -7,12 +7,17 @@ import (
"github.com/netsampler/goflow2/transport" "github.com/netsampler/goflow2/transport"
"io" "io"
"os" "os"
"os/signal"
"sync"
"syscall"
) )
type FileDriver struct { type FileDriver struct {
fileDestination string fileDestination string
w io.Writer w io.Writer
file *os.File file *os.File
lock *sync.RWMutex
q chan bool
} }
func (d *FileDriver) Prepare() error { func (d *FileDriver) Prepare() error {
@ -21,33 +26,74 @@ func (d *FileDriver) Prepare() error {
return nil return nil
} }
func (d *FileDriver) openFile() error {
file, err := os.OpenFile(d.fileDestination, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
if err != nil {
return err
}
d.file = file
d.w = d.file
return err
}
func (d *FileDriver) Init(context.Context) error { func (d *FileDriver) Init(context.Context) error {
d.q = make(chan bool, 1)
if d.fileDestination == "" { if d.fileDestination == "" {
d.w = os.Stdout d.w = os.Stdout
} else { } else {
var err error var err error
d.file, err = os.OpenFile(d.fileDestination, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
d.lock.Lock()
err = d.openFile()
d.lock.Unlock()
if err != nil { if err != nil {
return err return err
} }
d.w = d.file
c := make(chan os.Signal, 1)
signal.Notify(c, syscall.SIGHUP)
go func() {
for {
select {
case <-c:
d.lock.Lock()
d.file.Close()
d.openFile()
d.lock.Unlock()
// if there is an error, keeps using the old file
case <-d.q:
return
}
}
}()
} }
return nil return nil
} }
func (d *FileDriver) Send(key, data []byte) error { func (d *FileDriver) Send(key, data []byte) error {
fmt.Fprintln(d.w, string(data)) d.lock.RLock()
return nil w := d.w
d.lock.RUnlock()
_, err := fmt.Fprintln(w, string(data))
return err
} }
func (d *FileDriver) Close(context.Context) error { func (d *FileDriver) Close(context.Context) error {
if d.fileDestination != "" { if d.fileDestination != "" {
d.lock.Lock()
d.file.Close() d.file.Close()
d.lock.Unlock()
signal.Ignore(syscall.SIGHUP)
} }
close(d.q)
return nil return nil
} }
func init() { func init() {
d := &FileDriver{} d := &FileDriver{
lock: &sync.RWMutex{},
}
transport.RegisterTransportDriver("file", d) transport.RegisterTransportDriver("file", d)
} }