mirror of
https://github.com/netsampler/goflow2.git
synced 2024-05-06 15:54:52 +00:00
Merge pull request #4 from netsampler/feature/sighup
feature: implement SIGHUP for log rotation
This commit is contained in:
@ -7,12 +7,17 @@ import (
|
||||
"github.com/netsampler/goflow2/transport"
|
||||
"io"
|
||||
"os"
|
||||
"os/signal"
|
||||
"sync"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
type FileDriver struct {
|
||||
fileDestination string
|
||||
w io.Writer
|
||||
file *os.File
|
||||
lock *sync.RWMutex
|
||||
q chan bool
|
||||
}
|
||||
|
||||
func (d *FileDriver) Prepare() error {
|
||||
@ -21,33 +26,74 @@ func (d *FileDriver) Prepare() error {
|
||||
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 {
|
||||
d.q = make(chan bool, 1)
|
||||
|
||||
if d.fileDestination == "" {
|
||||
d.w = os.Stdout
|
||||
} else {
|
||||
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 {
|
||||
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
|
||||
}
|
||||
|
||||
func (d *FileDriver) Send(key, data []byte) error {
|
||||
fmt.Fprintln(d.w, string(data))
|
||||
return nil
|
||||
d.lock.RLock()
|
||||
w := d.w
|
||||
d.lock.RUnlock()
|
||||
_, err := fmt.Fprintln(w, string(data))
|
||||
return err
|
||||
}
|
||||
|
||||
func (d *FileDriver) Close(context.Context) error {
|
||||
if d.fileDestination != "" {
|
||||
d.lock.Lock()
|
||||
d.file.Close()
|
||||
d.lock.Unlock()
|
||||
signal.Ignore(syscall.SIGHUP)
|
||||
}
|
||||
close(d.q)
|
||||
return nil
|
||||
}
|
||||
|
||||
func init() {
|
||||
d := &FileDriver{}
|
||||
d := &FileDriver{
|
||||
lock: &sync.RWMutex{},
|
||||
}
|
||||
transport.RegisterTransportDriver("file", d)
|
||||
}
|
||||
|
Reference in New Issue
Block a user