diff --git a/os.go b/os.go index 1d0b538d..a273bec4 100644 --- a/os.go +++ b/os.go @@ -8,6 +8,7 @@ import ( "os/signal" "strings" "sync" + "syscall" "time" ) @@ -15,6 +16,10 @@ var ( GoPath = os.Getenv("GOPATH") ) +func init() { + initAFSIGHUPWatcher() +} + func TrapSignal(cb func()) { c := make(chan os.Signal, 1) signal.Notify(c, os.Interrupt) @@ -134,6 +139,7 @@ const autoFileOpenDuration = 1000 * time.Millisecond // Automatically closes and re-opens file for writing. // This is useful for using a log file with the logrotate tool. type AutoFile struct { + ID string Path string ticker *time.Ticker mtx sync.Mutex @@ -142,6 +148,7 @@ type AutoFile struct { func OpenAutoFile(path string) (af *AutoFile, err error) { af = &AutoFile{ + ID: RandStr(12) + ":" + path, Path: path, ticker: time.NewTicker(autoFileOpenDuration), } @@ -149,14 +156,14 @@ func OpenAutoFile(path string) (af *AutoFile, err error) { return } go af.processTicks() + autoFileWatchers.addAutoFile(af) return } func (af *AutoFile) Close() error { af.ticker.Stop() - af.mtx.Lock() err := af.closeFile() - af.mtx.Unlock() + autoFileWatchers.removeAutoFile(af) return err } @@ -166,13 +173,14 @@ func (af *AutoFile) processTicks() { if !ok { return // Done. } - af.mtx.Lock() af.closeFile() - af.mtx.Unlock() } } func (af *AutoFile) closeFile() (err error) { + af.mtx.Lock() + defer af.mtx.Unlock() + file := af.file if file == nil { return nil @@ -201,6 +209,56 @@ func (af *AutoFile) openFile() error { return nil } +//-------------------------------------------------------------------------------- + +var autoFileWatchers *afSIGHUPWatcher + +func initAFSIGHUPWatcher() { + autoFileWatchers = newAFSIGHUPWatcher() + + c := make(chan os.Signal, 1) + signal.Notify(c, syscall.SIGHUP) + + go func() { + for _ = range c { + autoFileWatchers.closeAll() + } + }() +} + +type afSIGHUPWatcher struct { + mtx sync.Mutex + autoFiles map[string]*AutoFile +} + +func newAFSIGHUPWatcher() *afSIGHUPWatcher { + return &afSIGHUPWatcher{ + autoFiles: make(map[string]*AutoFile, 10), + } +} + +func (afw *afSIGHUPWatcher) addAutoFile(af *AutoFile) { + afw.mtx.Lock() + afw.autoFiles[af.ID] = af + afw.mtx.Unlock() +} + +func (afw *afSIGHUPWatcher) removeAutoFile(af *AutoFile) { + afw.mtx.Lock() + delete(afw.autoFiles, af.ID) + afw.mtx.Unlock() +} + +func (afw *afSIGHUPWatcher) closeAll() { + afw.mtx.Lock() + for _, af := range afw.autoFiles { + af.closeFile() + } + afw.mtx.Unlock() +} + +//-------------------------------------------------------------------------------- + func Tempfile(prefix string) (*os.File, string) { file, err := ioutil.TempFile("", prefix) if err != nil { @@ -209,6 +267,8 @@ func Tempfile(prefix string) (*os.File, string) { return file, file.Name() } +//-------------------------------------------------------------------------------- + func Prompt(prompt string, defaultValue string) (string, error) { fmt.Print(prompt) reader := bufio.NewReader(os.Stdin) diff --git a/os_test.go b/os_test.go new file mode 100644 index 00000000..c0effdc2 --- /dev/null +++ b/os_test.go @@ -0,0 +1,64 @@ +package common + +import ( + "os" + "syscall" + "testing" +) + +func TestSIGHUP(t *testing.T) { + + // First, create an AutoFile writing to a tempfile dir + file, name := Tempfile("sighup_test") + err := file.Close() + if err != nil { + t.Fatalf("Error creating tempfile: %v", err) + } + // Here is the actual AutoFile + af, err := OpenAutoFile(name) + if err != nil { + t.Fatalf("Error creating autofile: %v", err) + } + + // Write to the file. + _, err = af.Write([]byte("Line 1\n")) + if err != nil { + t.Fatalf("Error writing to autofile: %v", err) + } + _, err = af.Write([]byte("Line 2\n")) + if err != nil { + t.Fatalf("Error writing to autofile: %v", err) + } + + // Send SIGHUP to self. + syscall.Kill(syscall.Getpid(), syscall.SIGHUP) + + // Move the file over + err = os.Rename(name, name+"_old") + if err != nil { + t.Fatalf("Error moving autofile: %v", err) + } + + // Write more to the file. + _, err = af.Write([]byte("Line 3\n")) + if err != nil { + t.Fatalf("Error writing to autofile: %v", err) + } + _, err = af.Write([]byte("Line 4\n")) + if err != nil { + t.Fatalf("Error writing to autofile: %v", err) + } + err = af.Close() + if err != nil { + t.Fatalf("Error closing autofile") + } + + // Both files should exist + if body := MustReadFile(name + "_old"); string(body) != "Line 1\nLine 2\n" { + t.Errorf("Unexpected body %s", body) + } + if body := MustReadFile(name); string(body) != "Line 3\nLine 4\n" { + t.Errorf("Unexpected body %s", body) + } + +}