store.go 1.9 KB

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