store.go 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. package main
  2. import (
  3. // "bufio"
  4. "encoding/gob"
  5. "io"
  6. "log"
  7. "os"
  8. "sync"
  9. )
  10. const saveQueueLength = 1000
  11. type URLStore struct {
  12. urls map[string]string
  13. mu sync.RWMutex
  14. save chan record
  15. }
  16. type record struct {
  17. Key, URL string
  18. }
  19. func NewURLStore(filename string) *URLStore {
  20. s := &URLStore{
  21. urls: make(map[string]string),
  22. save: make(chan record, saveQueueLength),
  23. }
  24. if err := s.load(filename); err != nil {
  25. log.Println("Error loading URLStore:", err)
  26. }
  27. go s.saveLoop(filename)
  28. return s
  29. }
  30. func (s *URLStore) Get(key string) string {
  31. s.mu.RLock()
  32. defer s.mu.RUnlock()
  33. return s.urls[key]
  34. }
  35. func (s *URLStore) Set(key, url string) bool {
  36. s.mu.Lock()
  37. defer s.mu.Unlock()
  38. if _, present := s.urls[key]; present {
  39. return false
  40. }
  41. s.urls[key] = url
  42. return true
  43. }
  44. func (s *URLStore) Count() int {
  45. s.mu.RLock()
  46. defer s.mu.RUnlock()
  47. return len(s.urls)
  48. }
  49. func (s *URLStore) Put(url string) string {
  50. for {
  51. key := genKey(s.Count())
  52. if ok := s.Set(key, url); ok {
  53. s.save <- record{key, url}
  54. return key
  55. }
  56. }
  57. panic("shouldn't get here")
  58. }
  59. func (s *URLStore) load(filename string) error {
  60. f, err := os.Open(filename)
  61. if err != nil {
  62. log.Println("Error opening URLStore:", err)
  63. return err
  64. }
  65. defer f.Close()
  66. // buffered reading:
  67. // b := bufio.NewReader(f)
  68. // d := gob.NewDecoder(b)
  69. d := gob.NewDecoder(f)
  70. for err == nil {
  71. var r record
  72. if err = d.Decode(&r); err == nil {
  73. s.Set(r.Key, r.URL)
  74. }
  75. }
  76. if err == io.EOF {
  77. return nil
  78. }
  79. // error occurred:
  80. log.Println("Error decoding URLStore:", err) // map hasn't been read correctly
  81. return err
  82. }
  83. func (s *URLStore) saveLoop(filename string) {
  84. f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
  85. if err != nil {
  86. log.Fatal("Error opening URLStore: ", err)
  87. }
  88. defer f.Close()
  89. e := gob.NewEncoder(f)
  90. // buffered encoding:
  91. // b := bufio.NewWriter(f)
  92. // e := gob.NewEncoder(b)
  93. // defer b.Flush()
  94. for {
  95. r := <-s.save // takes a record from the channel
  96. if err := e.Encode(r); err != nil {
  97. log.Println("Error saving to URLStore: ", err)
  98. }
  99. }
  100. }