memcached.go 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184
  1. package database
  2. import (
  3. "encoding/json"
  4. "errors"
  5. "log"
  6. "net"
  7. "strings"
  8. "time"
  9. "github.com/bradfitz/gomemcache/memcache"
  10. "github.com/gravitl/netmaker/servercfg"
  11. )
  12. var MemCachedDatabase *memcache.Client
  13. var MEMCACHED_FUNCTIONS = map[string]interface{}{
  14. INIT_DB: initMemcachedDatabase,
  15. CREATE_TABLE: memcachedCreateTable,
  16. INSERT: memcachedInsert,
  17. INSERT_PEER: memcachedInsertPeer,
  18. DELETE: memcachedDeleteRecord,
  19. DELETE_ALL: memcachedDeleteAllRecords,
  20. FETCH_ALL: memcachedFetchRecords,
  21. CLOSE_DB: memcachedCloseDB,
  22. }
  23. // utility function to make setting memcached servers easier
  24. func parseMemcachedAddresses(addresses string) string {
  25. addressesArr := strings.Split(addresses, ",")
  26. numAddresses := len(addressesArr)
  27. if numAddresses == 0 {
  28. return "127.0.0.1:11211"
  29. }
  30. newAddresses := ""
  31. log.SetFlags(log.Flags() &^ (log.Llongfile | log.Lshortfile))
  32. for _, address := range addressesArr {
  33. if isValidIp(address) {
  34. newAddresses += address
  35. if servercfg.GetVerbose() >= 2 {
  36. log.Println("adding " + address + " to memcached servers")
  37. }
  38. if address != addressesArr[numAddresses-1] {
  39. newAddresses += ","
  40. }
  41. }
  42. }
  43. return newAddresses
  44. }
  45. func initMemcachedDatabase() error {
  46. addresses := parseMemcachedAddresses(servercfg.GetMemcachedAddress())
  47. MemCachedDatabase = memcache.New(addresses)
  48. if MemCachedDatabase == nil {
  49. return errors.New("could not initialize memcached")
  50. }
  51. MemCachedDatabase.Timeout = time.Minute
  52. return nil
  53. }
  54. func memcachedCreateTable(tableName string) error {
  55. if currentTable, err := memcachedFetchRecords(tableName); (currentTable != nil && len(currentTable) >= 0) || err != nil {
  56. // return if it already exists
  57. return err
  58. } else {
  59. log.Println(currentTable)
  60. }
  61. table := make(map[string]string)
  62. newTable, err := json.Marshal(table)
  63. if err != nil {
  64. return err
  65. }
  66. err = MemCachedDatabase.Set(&memcache.Item{Key: tableName, Value: newTable, Expiration: 0})
  67. if err != nil {
  68. return err
  69. }
  70. return nil
  71. }
  72. func memcachedInsert(key string, value string, tableName string) error {
  73. if key != "" && value != "" && IsJSONString(value) {
  74. preData, err := MemCachedDatabase.Get(tableName)
  75. if err != nil {
  76. return err
  77. }
  78. var preDataMap map[string]string
  79. if err := json.Unmarshal(preData.Value, &preDataMap); err != nil {
  80. return err
  81. }
  82. preDataMap[key] = value
  83. postData, err := json.Marshal(&preDataMap)
  84. if err != nil {
  85. return err
  86. }
  87. err = MemCachedDatabase.Replace(&memcache.Item{Key: tableName, Value: postData, Expiration: 0})
  88. if err != nil {
  89. return err
  90. }
  91. return nil
  92. } else {
  93. return errors.New("invalid insert " + key + " : " + value)
  94. }
  95. }
  96. func memcachedInsertPeer(key string, value string) error {
  97. if key != "" && value != "" && IsJSONString(value) {
  98. if err := memcachedInsert(key, value, PEERS_TABLE_NAME); err != nil {
  99. return err
  100. }
  101. return nil
  102. } else {
  103. return errors.New("invalid peer insert " + key + " : " + value)
  104. }
  105. }
  106. func memcachedDeleteRecord(tableName string, key string) error {
  107. if key != "" {
  108. preData, err := MemCachedDatabase.Get(tableName)
  109. if err != nil {
  110. return err
  111. }
  112. var preDataMap map[string]string
  113. if err := json.Unmarshal(preData.Value, &preDataMap); err != nil {
  114. return err
  115. }
  116. delete(preDataMap, key)
  117. postData, err := json.Marshal(&preDataMap)
  118. if err != nil {
  119. return err
  120. }
  121. err = MemCachedDatabase.Set(&memcache.Item{Key: tableName, Value: postData, Expiration: 0})
  122. if err != nil {
  123. return err
  124. }
  125. return nil
  126. } else {
  127. return errors.New("invalid delete, key is required")
  128. }
  129. }
  130. func memcachedDeleteAllRecords(tableName string) error {
  131. err := MemCachedDatabase.Delete(tableName)
  132. if err != nil {
  133. return err
  134. }
  135. err = memcachedCreateTable(tableName)
  136. if err != nil {
  137. return err
  138. }
  139. return nil
  140. }
  141. // func memcachedFetchRecord(tableName string, key string) (string, error) {
  142. // results, err := memcachedFetchRecords(tableName)
  143. // if err != nil {
  144. // return "", err
  145. // }
  146. // if results[key] == "" {
  147. // return "", errors.New(NO_RECORD)
  148. // }
  149. // return results[key], nil
  150. // }
  151. func memcachedFetchRecords(tableName string) (map[string]string, error) {
  152. var records map[string]string
  153. item, err := MemCachedDatabase.Get(tableName)
  154. if err != nil {
  155. return records, err
  156. }
  157. if err = json.Unmarshal(item.Value, &records); err != nil {
  158. return nil, err
  159. }
  160. return records, nil
  161. }
  162. func memcachedCloseDB() {
  163. // no op for this library..
  164. }
  165. func isValidIp(ipAddr string) bool {
  166. return net.ParseIP(ipAddr) == nil
  167. }