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.

225 lines
4.5 KiB

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