api.go 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534
  1. package pudge
  2. import (
  3. "bytes"
  4. "encoding/gob"
  5. "os"
  6. )
  7. // DefaultConfig is default config
  8. var DefaultConfig = &Config{
  9. FileMode: 0666,
  10. DirMode: 0777,
  11. SyncInterval: 0,
  12. StoreMode: 0}
  13. // Open return db object if it opened.
  14. // Create new db if not exist.
  15. // Read db to obj if exist.
  16. // Or error if any.
  17. // Default Config (if nil): &Config{FileMode: 0666, DirMode: 0777, SyncInterval: 0}
  18. func Open(f string, cfg *Config) (*Db, error) {
  19. if cfg == nil {
  20. cfg = DefaultConfig
  21. }
  22. dbs.RLock()
  23. db, ok := dbs.dbs[f]
  24. if ok {
  25. dbs.RUnlock()
  26. return db, nil
  27. }
  28. dbs.RUnlock()
  29. dbs.Lock()
  30. db, err := newDb(f, cfg)
  31. //log.Println("n", db.name, db.config.StoreMode)
  32. if err == nil {
  33. dbs.dbs[f] = db
  34. }
  35. dbs.Unlock()
  36. return db, err
  37. }
  38. // Set store any key value to db
  39. func (db *Db) Set(key, value interface{}) error {
  40. db.Lock()
  41. defer db.Unlock()
  42. k, err := KeyToBinary(key)
  43. if err != nil {
  44. return err
  45. }
  46. v, err := ValToBinary(value)
  47. if err != nil {
  48. return err
  49. }
  50. //log.Println("Set:", k, v)
  51. oldCmd, exists := db.vals[string(k)]
  52. //fmt.Println("StoreMode", db.config.StoreMode)
  53. if db.storemode == 2 {
  54. cmd := &Cmd{}
  55. cmd.Size = uint32(len(v))
  56. cmd.Val = make([]byte, len(v))
  57. copy(cmd.Val, v)
  58. db.vals[string(k)] = cmd
  59. } else {
  60. cmd, err := writeKeyVal(db.fk, db.fv, k, v, exists, oldCmd)
  61. if err != nil {
  62. return err
  63. }
  64. db.vals[string(k)] = cmd
  65. }
  66. if !exists {
  67. db.appendKey(k)
  68. }
  69. return err
  70. }
  71. // Get return value by key
  72. // Return error if any.
  73. func (db *Db) Get(key, value interface{}) error {
  74. db.RLock()
  75. defer db.RUnlock()
  76. k, err := KeyToBinary(key)
  77. if err != nil {
  78. return err
  79. }
  80. if val, ok := db.vals[string(k)]; ok {
  81. switch value.(type) {
  82. case *[]byte:
  83. b := make([]byte, val.Size)
  84. if db.storemode == 2 {
  85. copy(b, val.Val)
  86. } else {
  87. _, err := db.fv.ReadAt(b, int64(val.Seek))
  88. if err != nil {
  89. return err
  90. }
  91. }
  92. *value.(*[]byte) = b
  93. return nil
  94. default:
  95. buf := new(bytes.Buffer)
  96. b := make([]byte, val.Size)
  97. if db.storemode == 2 {
  98. //fmt.Println(val)
  99. copy(b, val.Val)
  100. } else {
  101. _, err := db.fv.ReadAt(b, int64(val.Seek))
  102. if err != nil {
  103. return err
  104. }
  105. }
  106. buf.Write(b)
  107. err = gob.NewDecoder(buf).Decode(value)
  108. return err
  109. }
  110. }
  111. return ErrKeyNotFound
  112. }
  113. // Close - sync & close files.
  114. // Return error if any.
  115. func (db *Db) Close() error {
  116. if db.cancelSyncer != nil {
  117. db.cancelSyncer()
  118. }
  119. db.Lock()
  120. defer db.Unlock()
  121. if db.storemode == 2 && db.name != "" {
  122. db.sort()
  123. keys := make([][]byte, len(db.keys))
  124. copy(keys, db.keys)
  125. db.storemode = 0
  126. for _, k := range keys {
  127. if val, ok := db.vals[string(k)]; ok {
  128. writeKeyVal(db.fk, db.fv, k, val.Val, false, nil)
  129. }
  130. }
  131. }
  132. if db.fk != nil {
  133. err := db.fk.Sync()
  134. if err != nil {
  135. return err
  136. }
  137. err = db.fk.Close()
  138. if err != nil {
  139. return err
  140. }
  141. }
  142. if db.fv != nil {
  143. err := db.fv.Sync()
  144. if err != nil {
  145. return err
  146. }
  147. err = db.fv.Close()
  148. if err != nil {
  149. return err
  150. }
  151. }
  152. dbs.Lock()
  153. delete(dbs.dbs, db.name)
  154. dbs.Unlock()
  155. return nil
  156. }
  157. // CloseAll - close all opened Db
  158. func CloseAll() (err error) {
  159. dbs.Lock()
  160. stores := dbs.dbs
  161. dbs.Unlock()
  162. for _, db := range stores {
  163. err = db.Close()
  164. if err != nil {
  165. break
  166. }
  167. }
  168. return err
  169. }
  170. // DeleteFile close and delete file
  171. func (db *Db) DeleteFile() error {
  172. return DeleteFile(db.name)
  173. }
  174. // DeleteFile close db and delete file
  175. func DeleteFile(file string) error {
  176. if file == "" {
  177. return nil
  178. }
  179. dbs.Lock()
  180. db, ok := dbs.dbs[file]
  181. if ok {
  182. dbs.Unlock()
  183. err := db.Close()
  184. if err != nil {
  185. return err
  186. }
  187. } else {
  188. dbs.Unlock()
  189. }
  190. err := os.Remove(file)
  191. if err != nil {
  192. return err
  193. }
  194. err = os.Remove(file + ".idx")
  195. return err
  196. }
  197. // Has return true if key exists.
  198. // Return error if any.
  199. func (db *Db) Has(key interface{}) (bool, error) {
  200. db.RLock()
  201. defer db.RUnlock()
  202. k, err := KeyToBinary(key)
  203. if err != nil {
  204. return false, err
  205. }
  206. _, has := db.vals[string(k)]
  207. return has, nil
  208. }
  209. // FileSize returns the total size of the disk storage used by the DB.
  210. func (db *Db) FileSize() (int64, error) {
  211. db.RLock()
  212. defer db.RUnlock()
  213. var err error
  214. is, err := db.fk.Stat()
  215. if err != nil {
  216. return -1, err
  217. }
  218. ds, err := db.fv.Stat()
  219. if err != nil {
  220. return -1, err
  221. }
  222. return is.Size() + ds.Size(), nil
  223. }
  224. // Count returns the number of items in the Db.
  225. func (db *Db) Count() (int, error) {
  226. db.RLock()
  227. defer db.RUnlock()
  228. return len(db.keys), nil
  229. }
  230. // Delete remove key
  231. // Returns error if key not found
  232. func (db *Db) Delete(key interface{}) error {
  233. db.Lock()
  234. defer db.Unlock()
  235. k, err := KeyToBinary(key)
  236. if err != nil {
  237. return err
  238. }
  239. if _, ok := db.vals[string(k)]; ok {
  240. delete(db.vals, string(k))
  241. db.deleteFromKeys(k)
  242. writeKey(db.fk, 1, 0, 0, k, -1)
  243. return nil
  244. }
  245. return ErrKeyNotFound
  246. }
  247. // KeysByPrefix return keys with prefix
  248. // in ascending or descending order (false - descending,true - ascending)
  249. // if limit == 0 return all keys
  250. // if offset > 0 - skip offset records
  251. // If from not nil - return keys after from (from not included)
  252. func (db *Db) KeysByPrefix(prefix []byte, limit, offset int, asc bool) ([][]byte, error) {
  253. //log.Println("KeysByPrefix")
  254. db.RLock()
  255. defer db.RUnlock()
  256. // resulting array
  257. arr := make([][]byte, 0, 0)
  258. found := db.foundPref(prefix, asc)
  259. if found >= len(db.keys) || !startFrom(db.keys[found], prefix) {
  260. //not found
  261. return arr, ErrKeyNotFound
  262. }
  263. start, end := checkInterval(found, limit, offset, 0, len(db.keys), asc)
  264. if start < 0 || start >= len(db.keys) {
  265. return arr, nil
  266. }
  267. if asc {
  268. for i := start; i <= end; i++ {
  269. if !startFrom(db.keys[i], prefix) {
  270. break
  271. }
  272. arr = append(arr, db.keys[i])
  273. }
  274. } else {
  275. for i := start; i >= end; i-- {
  276. if !startFrom(db.keys[i], prefix) {
  277. break
  278. }
  279. arr = append(arr, db.keys[i])
  280. }
  281. }
  282. return arr, nil
  283. }
  284. // Keys return keys in ascending or descending order (false - descending,true - ascending)
  285. // if limit == 0 return all keys
  286. // if offset > 0 - skip offset records
  287. // If from not nil - return keys after from (from not included)
  288. func (db *Db) Keys(from interface{}, limit, offset int, asc bool) ([][]byte, error) {
  289. // resulting array
  290. //log.Println("pudge", from, from == nil)
  291. arr := make([][]byte, 0, 0)
  292. excludeFrom := 0
  293. if from != nil {
  294. excludeFrom = 1
  295. k, err := KeyToBinary(from)
  296. //log.Println(bytes.Equal(k[len(k)-1:], []byte("*")))
  297. if err != nil {
  298. return arr, err
  299. }
  300. if len(k) > 1 && bytes.Equal(k[len(k)-1:], []byte("*")) {
  301. byteOrStr := false
  302. switch from.(type) {
  303. case []byte:
  304. byteOrStr = true
  305. case string:
  306. byteOrStr = true
  307. }
  308. if byteOrStr {
  309. prefix := make([]byte, len(k)-1)
  310. copy(prefix, k)
  311. return db.KeysByPrefix(prefix, limit, offset, asc)
  312. }
  313. }
  314. }
  315. db.RLock()
  316. defer db.RUnlock()
  317. find, _ := db.findKey(from, asc)
  318. start, end := checkInterval(find, limit, offset, excludeFrom, len(db.keys), asc)
  319. if start < 0 || start >= len(db.keys) {
  320. return arr, nil
  321. }
  322. if asc {
  323. for i := start; i <= end; i++ {
  324. arr = append(arr, db.keys[i])
  325. }
  326. } else {
  327. for i := start; i >= end; i-- {
  328. arr = append(arr, db.keys[i])
  329. }
  330. }
  331. return arr, nil
  332. }
  333. // Counter return int64 incremented on incr
  334. func (db *Db) Counter(key interface{}, incr int) (int64, error) {
  335. mutex.Lock()
  336. var counter int64
  337. err := db.Get(key, &counter)
  338. if err != nil && err != ErrKeyNotFound {
  339. return -1, err
  340. }
  341. //mutex.Lock()
  342. counter = counter + int64(incr)
  343. //mutex.Unlock()
  344. err = db.Set(key, counter)
  345. mutex.Unlock()
  346. return counter, err
  347. }
  348. // Set store any key value to db with opening if needed
  349. func Set(f string, key, value interface{}) error {
  350. db, err := Open(f, nil)
  351. if err != nil {
  352. return err
  353. }
  354. return db.Set(key, value)
  355. }
  356. // Sets store vals and keys
  357. // Use it for mass insertion
  358. // every pair must contain key and value
  359. func Sets(file string, pairs []interface{}) (err error) {
  360. db, err := Open(file, nil)
  361. if err != nil {
  362. return err
  363. }
  364. for i := range pairs {
  365. if i%2 != 0 {
  366. // on odd - append val and store key
  367. if pairs[i] == nil || pairs[i-1] == nil {
  368. break
  369. }
  370. err = db.Set(pairs[i-1], pairs[i])
  371. if err != nil {
  372. break
  373. }
  374. }
  375. }
  376. return err
  377. }
  378. // Get return value by key with opening if needed
  379. // Return error if any.
  380. func Get(f string, key, value interface{}) error {
  381. db, err := Open(f, nil)
  382. if err != nil {
  383. return err
  384. }
  385. return db.Get(key, value)
  386. }
  387. // Gets return key/value pairs in random order
  388. // result contains key and value
  389. // Gets not return error if key not found
  390. // If no keys found return empty result
  391. func Gets(file string, keys []interface{}) (result [][]byte) {
  392. db, err := Open(file, nil)
  393. if err != nil {
  394. return nil
  395. }
  396. for _, key := range keys {
  397. var v []byte
  398. err := db.Get(key, &v)
  399. if err == nil {
  400. k, err := KeyToBinary(key)
  401. if err == nil {
  402. val, err := ValToBinary(v)
  403. if err == nil {
  404. result = append(result, k)
  405. result = append(result, val)
  406. }
  407. }
  408. }
  409. }
  410. return result
  411. }
  412. // Counter return int64 incremented on incr with lazy open
  413. func Counter(f string, key interface{}, incr int) (int64, error) {
  414. db, err := Open(f, nil)
  415. if err != nil {
  416. return 0, err
  417. }
  418. return db.Counter(key, incr)
  419. }
  420. // Delete remove key
  421. // Returns error if key not found
  422. func Delete(f string, key interface{}) error {
  423. db, err := Open(f, nil)
  424. if err != nil {
  425. return err
  426. }
  427. return db.Delete(key)
  428. }
  429. // Keys return keys in ascending or descending order (false - descending,true - ascending)
  430. // if limit == 0 return all keys
  431. // if offset > 0 - skip offset records
  432. // If from not nil - return keys after from (from not included)
  433. func Keys(f string, from interface{}, limit, offset int, asc bool) ([][]byte, error) {
  434. db, err := Open(f, nil)
  435. if err != nil {
  436. return nil, err
  437. }
  438. return db.Keys(from, limit, offset, asc)
  439. }
  440. // Has return true if key exists.
  441. // Return error if any.
  442. func Has(f string, key interface{}) (bool, error) {
  443. db, err := Open(f, nil)
  444. if err != nil {
  445. return false, err
  446. }
  447. return db.Has(key)
  448. }
  449. // Count returns the number of items in the Db.
  450. func Count(f string) (int, error) {
  451. db, err := Open(f, nil)
  452. if err != nil {
  453. return -1, err
  454. }
  455. return db.Count()
  456. }
  457. // Close - sync & close files.
  458. // Return error if any.
  459. func Close(f string) error {
  460. db, err := Open(f, nil)
  461. if err != nil {
  462. return err
  463. }
  464. return db.Close()
  465. }
  466. // BackupAll - backup all opened Db
  467. // if dir not set it will be backup
  468. // delete old backup file before run
  469. // ignore all errors
  470. func BackupAll(dir string) (err error) {
  471. if dir == "" {
  472. dir = "backup"
  473. }
  474. dbs.Lock()
  475. stores := dbs.dbs
  476. dbs.Unlock()
  477. //tmp := make(map[string]string)
  478. for _, db := range stores {
  479. backup := dir + "/" + db.name
  480. DeleteFile(backup)
  481. keys, err := db.Keys(nil, 0, 0, true)
  482. if err == nil {
  483. for _, k := range keys {
  484. var b []byte
  485. db.Get(k, &b)
  486. Set(backup, k, b)
  487. }
  488. }
  489. Close(backup)
  490. }
  491. return err
  492. }