context.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183
  1. package context
  2. import (
  3. "bytes"
  4. "encoding/binary"
  5. "io"
  6. "io/ioutil"
  7. "os"
  8. "path/filepath"
  9. "sync"
  10. "time"
  11. "github.com/syndtr/goleveldb/leveldb"
  12. "github.com/syndtr/goleveldb/leveldb/opt"
  13. )
  14. const (
  15. routesDbFilename = "routes.db"
  16. idLogFilename = "id"
  17. )
  18. // Route is the value part of a shortcut.
  19. type Route struct {
  20. URL string `json:"url"`
  21. Time time.Time `json:"time"`
  22. }
  23. // Serialize this Route into the given Writer.
  24. func (o *Route) write(w io.Writer) error {
  25. if err := binary.Write(w, binary.LittleEndian, o.Time.UnixNano()); err != nil {
  26. return err
  27. }
  28. if _, err := w.Write([]byte(o.URL)); err != nil {
  29. return err
  30. }
  31. return nil
  32. }
  33. // Deserialize this Route from the given Reader.
  34. func (o *Route) read(r io.Reader) error {
  35. var t int64
  36. if err := binary.Read(r, binary.LittleEndian, &t); err != nil {
  37. return err
  38. }
  39. b, err := ioutil.ReadAll(r)
  40. if err != nil {
  41. return err
  42. }
  43. o.URL = string(b)
  44. o.Time = time.Unix(0, t)
  45. return nil
  46. }
  47. // Context provides access to the data store.
  48. type Context struct {
  49. path string
  50. db *leveldb.DB
  51. lck sync.Mutex
  52. id uint64
  53. }
  54. // Commit the given ID to the data store.
  55. func commit(filename string, id uint64) error {
  56. w, err := os.Create(filename)
  57. if err != nil {
  58. return err
  59. }
  60. defer w.Close()
  61. if err := binary.Write(w, binary.LittleEndian, id); err != nil {
  62. return err
  63. }
  64. return w.Sync()
  65. }
  66. // Load the current ID from the data store.
  67. func load(filename string) (uint64, error) {
  68. if _, err := os.Stat(filename); err != nil {
  69. return 0, commit(filename, 0)
  70. }
  71. r, err := os.Open(filename)
  72. if err != nil {
  73. return 0, err
  74. }
  75. defer r.Close()
  76. var id uint64
  77. if err := binary.Read(r, binary.LittleEndian, &id); err != nil {
  78. return 0, err
  79. }
  80. return id, nil
  81. }
  82. // Open the context using path as the data store location.
  83. func Open(path string) (*Context, error) {
  84. if _, err := os.Stat(path); err != nil {
  85. if err := os.MkdirAll(path, os.ModePerm); err != nil {
  86. return nil, err
  87. }
  88. }
  89. // open the database
  90. db, err := leveldb.OpenFile(filepath.Join(path, routesDbFilename), nil)
  91. if err != nil {
  92. return nil, err
  93. }
  94. id, err := load(filepath.Join(path, idLogFilename))
  95. if err != nil {
  96. return nil, err
  97. }
  98. return &Context{
  99. path: path,
  100. db: db,
  101. id: id,
  102. }, nil
  103. }
  104. // Get retreives a shortcut from the data store.
  105. func (c *Context) Get(name string) (*Route, error) {
  106. val, err := c.db.Get([]byte(name), nil)
  107. if err != nil {
  108. return nil, err
  109. }
  110. rt := &Route{}
  111. if err := rt.read(bytes.NewBuffer(val)); err != nil {
  112. return nil, err
  113. }
  114. return rt, nil
  115. }
  116. // Put stores a new shortcut in the data store.
  117. func (c *Context) Put(key string, rt *Route) error {
  118. var buf bytes.Buffer
  119. if err := rt.write(&buf); err != nil {
  120. return err
  121. }
  122. return c.db.Put([]byte(key), buf.Bytes(), &opt.WriteOptions{Sync: true})
  123. }
  124. // Del removes an existing shortcut from the data store.
  125. func (c *Context) Del(key string) error {
  126. return c.db.Delete([]byte(key), &opt.WriteOptions{Sync: true})
  127. }
  128. func (c *Context) commit(id uint64) error {
  129. w, err := os.Create(filepath.Join(c.path, idLogFilename))
  130. if err != nil {
  131. return err
  132. }
  133. defer w.Close()
  134. if err := binary.Write(w, binary.LittleEndian, id); err != nil {
  135. return err
  136. }
  137. return w.Sync()
  138. }
  139. // NextID generates the next numeric ID to be used for an auto-named shortcut.
  140. func (c *Context) NextID() (uint64, error) {
  141. c.lck.Lock()
  142. defer c.lck.Unlock()
  143. c.id++
  144. if err := commit(filepath.Join(c.path, idLogFilename), c.id); err != nil {
  145. return 0, err
  146. }
  147. return c.id, nil
  148. }