|
|
- package common
-
- import (
- "fmt"
- "io/ioutil"
- "os"
- "os/signal"
- "sync"
- "time"
- )
-
- var (
- GoPath = os.Getenv("GOPATH")
- )
-
- func TrapSignal(cb func()) {
- c := make(chan os.Signal, 1)
- signal.Notify(c, os.Interrupt)
- signal.Notify(c, os.Kill)
- go func() {
- for sig := range c {
- fmt.Printf("captured %v, exiting...\n", sig)
- if cb != nil {
- cb()
- }
- os.Exit(1)
- }
- }()
- select {}
- }
-
- func Exit(s string) {
- fmt.Printf(s + "\n")
- os.Exit(1)
- }
-
- func EnsureDir(dir string) error {
- if _, err := os.Stat(dir); os.IsNotExist(err) {
- err := os.MkdirAll(dir, 0700)
- if err != nil {
- return fmt.Errorf("Could not create directory %v. %v", dir, err)
- }
- }
- return nil
- }
-
- func FileExists(filePath string) bool {
- _, err := os.Stat(filePath)
- return !os.IsNotExist(err)
- }
-
- func ReadFile(filePath string) ([]byte, error) {
- return ioutil.ReadFile(filePath)
- }
-
- func MustReadFile(filePath string) []byte {
- fileBytes, err := ioutil.ReadFile(filePath)
- if err != nil {
- Exit(Fmt("MustReadFile failed: %v", err))
- return nil
- }
- return fileBytes
- }
-
- func WriteFile(filePath string, contents []byte) error {
- err := ioutil.WriteFile(filePath, contents, 0600)
- if err != nil {
- return err
- }
- fmt.Printf("File written to %v.\n", filePath)
- return nil
- }
-
- func MustWriteFile(filePath string, contents []byte) {
- err := WriteFile(filePath, contents)
- if err != nil {
- Exit(Fmt("MustWriteFile failed: %v", err))
- }
- }
-
- // Writes to newBytes to filePath.
- // Guaranteed not to lose *both* oldBytes and newBytes,
- // (assuming that the OS is perfect)
- func WriteFileAtomic(filePath string, newBytes []byte) error {
- // If a file already exists there, copy to filePath+".bak" (overwrite anything)
- if _, err := os.Stat(filePath); !os.IsNotExist(err) {
- fileBytes, err := ioutil.ReadFile(filePath)
- if err != nil {
- return fmt.Errorf("Could not read file %v. %v", filePath, err)
- }
- err = ioutil.WriteFile(filePath+".bak", fileBytes, 0600)
- if err != nil {
- return fmt.Errorf("Could not write file %v. %v", filePath+".bak", err)
- }
- }
- // Write newBytes to filePath.new
- err := ioutil.WriteFile(filePath+".new", newBytes, 0600)
- if err != nil {
- return fmt.Errorf("Could not write file %v. %v", filePath+".new", err)
- }
- // Move filePath.new to filePath
- err = os.Rename(filePath+".new", filePath)
- return err
- }
-
- //--------------------------------------------------------------------------------
-
- /* AutoFile usage
-
- // Create/Append to ./autofile_test
- af, err := OpenAutoFile("autofile_test")
- if err != nil {
- panic(err)
- }
-
- // Stream of writes.
- // During this time, the file may be moved e.g. by logRotate.
- for i := 0; i < 60; i++ {
- af.Write([]byte(Fmt("LOOP(%v)", i)))
- time.Sleep(time.Second)
- }
-
- // Close the AutoFile
- err = af.Close()
- if err != nil {
- panic(err)
- }
- */
-
- 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 {
- Path string
- ticker *time.Ticker
- mtx sync.Mutex
- file *os.File
- }
-
- func OpenAutoFile(path string) (af *AutoFile, err error) {
- af = &AutoFile{
- Path: path,
- ticker: time.NewTicker(autoFileOpenDuration),
- }
- if err = af.openFile(); err != nil {
- return
- }
- go af.processTicks()
- return
- }
-
- func (af *AutoFile) Close() error {
- af.ticker.Stop()
- af.mtx.Lock()
- err := af.closeFile()
- af.mtx.Unlock()
- return err
- }
-
- func (af *AutoFile) processTicks() {
- for {
- _, ok := <-af.ticker.C
- if !ok {
- return // Done.
- }
- af.mtx.Lock()
- af.closeFile()
- af.mtx.Unlock()
- }
- }
-
- func (af *AutoFile) closeFile() (err error) {
- file := af.file
- if file == nil {
- return nil
- }
- af.file = nil
- return file.Close()
- }
-
- func (af *AutoFile) Write(b []byte) (n int, err error) {
- af.mtx.Lock()
- defer af.mtx.Unlock()
- if af.file == nil {
- if err = af.openFile(); err != nil {
- return
- }
- }
- return af.file.Write(b)
- }
-
- func (af *AutoFile) openFile() error {
- file, err := os.OpenFile(af.Path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
- if err != nil {
- return err
- }
- af.file = file
- return nil
- }
|