You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

200 lines
4.0 KiB

10 years ago
  1. package common
  2. import (
  3. "fmt"
  4. "io/ioutil"
  5. "os"
  6. "os/signal"
  7. "sync"
  8. "time"
  9. )
  10. var (
  11. GoPath = os.Getenv("GOPATH")
  12. )
  13. func TrapSignal(cb func()) {
  14. c := make(chan os.Signal, 1)
  15. signal.Notify(c, os.Interrupt)
  16. signal.Notify(c, os.Kill)
  17. go func() {
  18. for sig := range c {
  19. fmt.Printf("captured %v, exiting...\n", sig)
  20. if cb != nil {
  21. cb()
  22. }
  23. os.Exit(1)
  24. }
  25. }()
  26. select {}
  27. }
  28. func Exit(s string) {
  29. fmt.Printf(s + "\n")
  30. os.Exit(1)
  31. }
  32. func EnsureDir(dir string) error {
  33. if _, err := os.Stat(dir); os.IsNotExist(err) {
  34. err := os.MkdirAll(dir, 0700)
  35. if err != nil {
  36. return fmt.Errorf("Could not create directory %v. %v", dir, err)
  37. }
  38. }
  39. return nil
  40. }
  41. func FileExists(filePath string) bool {
  42. _, err := os.Stat(filePath)
  43. return !os.IsNotExist(err)
  44. }
  45. func ReadFile(filePath string) ([]byte, error) {
  46. return ioutil.ReadFile(filePath)
  47. }
  48. func MustReadFile(filePath string) []byte {
  49. fileBytes, err := ioutil.ReadFile(filePath)
  50. if err != nil {
  51. Exit(Fmt("MustReadFile failed: %v", err))
  52. return nil
  53. }
  54. return fileBytes
  55. }
  56. func WriteFile(filePath string, contents []byte) error {
  57. err := ioutil.WriteFile(filePath, contents, 0600)
  58. if err != nil {
  59. return err
  60. }
  61. fmt.Printf("File written to %v.\n", filePath)
  62. return nil
  63. }
  64. func MustWriteFile(filePath string, contents []byte) {
  65. err := WriteFile(filePath, contents)
  66. if err != nil {
  67. Exit(Fmt("MustWriteFile failed: %v", err))
  68. }
  69. }
  70. // Writes to newBytes to filePath.
  71. // Guaranteed not to lose *both* oldBytes and newBytes,
  72. // (assuming that the OS is perfect)
  73. func WriteFileAtomic(filePath string, newBytes []byte) error {
  74. // If a file already exists there, copy to filePath+".bak" (overwrite anything)
  75. if _, err := os.Stat(filePath); !os.IsNotExist(err) {
  76. fileBytes, err := ioutil.ReadFile(filePath)
  77. if err != nil {
  78. return fmt.Errorf("Could not read file %v. %v", filePath, err)
  79. }
  80. err = ioutil.WriteFile(filePath+".bak", fileBytes, 0600)
  81. if err != nil {
  82. return fmt.Errorf("Could not write file %v. %v", filePath+".bak", err)
  83. }
  84. }
  85. // Write newBytes to filePath.new
  86. err := ioutil.WriteFile(filePath+".new", newBytes, 0600)
  87. if err != nil {
  88. return fmt.Errorf("Could not write file %v. %v", filePath+".new", err)
  89. }
  90. // Move filePath.new to filePath
  91. err = os.Rename(filePath+".new", filePath)
  92. return err
  93. }
  94. //--------------------------------------------------------------------------------
  95. /* AutoFile usage
  96. // Create/Append to ./autofile_test
  97. af, err := OpenAutoFile("autofile_test")
  98. if err != nil {
  99. panic(err)
  100. }
  101. // Stream of writes.
  102. // During this time, the file may be moved e.g. by logRotate.
  103. for i := 0; i < 60; i++ {
  104. af.Write([]byte(Fmt("LOOP(%v)", i)))
  105. time.Sleep(time.Second)
  106. }
  107. // Close the AutoFile
  108. err = af.Close()
  109. if err != nil {
  110. panic(err)
  111. }
  112. */
  113. const autoFileOpenDuration = 1000 * time.Millisecond
  114. // Automatically closes and re-opens file for writing.
  115. // This is useful for using a log file with the logrotate tool.
  116. type AutoFile struct {
  117. Path string
  118. ticker *time.Ticker
  119. mtx sync.Mutex
  120. file *os.File
  121. }
  122. func OpenAutoFile(path string) (af *AutoFile, err error) {
  123. af = &AutoFile{
  124. Path: path,
  125. ticker: time.NewTicker(autoFileOpenDuration),
  126. }
  127. if err = af.openFile(); err != nil {
  128. return
  129. }
  130. go af.processTicks()
  131. return
  132. }
  133. func (af *AutoFile) Close() error {
  134. af.ticker.Stop()
  135. af.mtx.Lock()
  136. err := af.closeFile()
  137. af.mtx.Unlock()
  138. return err
  139. }
  140. func (af *AutoFile) processTicks() {
  141. for {
  142. _, ok := <-af.ticker.C
  143. if !ok {
  144. return // Done.
  145. }
  146. af.mtx.Lock()
  147. af.closeFile()
  148. af.mtx.Unlock()
  149. }
  150. }
  151. func (af *AutoFile) closeFile() (err error) {
  152. file := af.file
  153. if file == nil {
  154. return nil
  155. }
  156. af.file = nil
  157. return file.Close()
  158. }
  159. func (af *AutoFile) Write(b []byte) (n int, err error) {
  160. af.mtx.Lock()
  161. defer af.mtx.Unlock()
  162. if af.file == nil {
  163. if err = af.openFile(); err != nil {
  164. return
  165. }
  166. }
  167. return af.file.Write(b)
  168. }
  169. func (af *AutoFile) openFile() error {
  170. file, err := os.OpenFile(af.Path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
  171. if err != nil {
  172. return err
  173. }
  174. af.file = file
  175. return nil
  176. }