pudge.go 9.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421
  1. package pudge
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/binary"
  6. "encoding/gob"
  7. "errors"
  8. "io/ioutil"
  9. "os"
  10. "path/filepath"
  11. "sort"
  12. "sync"
  13. "time"
  14. )
  15. var (
  16. dbs struct {
  17. sync.RWMutex
  18. dbs map[string]*Db
  19. }
  20. // ErrKeyNotFound - key not found
  21. ErrKeyNotFound = errors.New("Error: key not found")
  22. mutex = &sync.RWMutex{}
  23. )
  24. // Db represent database
  25. type Db struct {
  26. sync.RWMutex
  27. name string
  28. fk *os.File
  29. fv *os.File
  30. keys [][]byte
  31. vals map[string]*Cmd
  32. cancelSyncer context.CancelFunc
  33. storemode int
  34. }
  35. // Cmd represent keys and vals addresses
  36. type Cmd struct {
  37. Seek uint32
  38. Size uint32
  39. KeySeek uint32
  40. Val []byte
  41. }
  42. // Config fo db
  43. // Default FileMode = 0666
  44. // Default DirMode = 0777
  45. // Default SyncInterval = 0 sec, 0 - disable sync (os will sync, typically 30 sec or so)
  46. // If StroreMode==2 && file == "" - pure inmemory mode
  47. type Config struct {
  48. FileMode int // 0666
  49. DirMode int // 0777
  50. SyncInterval int // in seconds
  51. StoreMode int // 0 - file first, 2 - memory first(with persist on close), 2 - with empty file - memory without persist
  52. }
  53. func init() {
  54. dbs.dbs = make(map[string]*Db)
  55. }
  56. func newDb(f string, cfg *Config) (*Db, error) {
  57. //fmt.Println("newdb2:", f, cfg.StoreMode)
  58. var err error
  59. // create
  60. db := new(Db)
  61. db.Lock()
  62. defer db.Unlock()
  63. // init
  64. db.name = f
  65. db.keys = make([][]byte, 0)
  66. db.vals = make(map[string]*Cmd)
  67. db.storemode = cfg.StoreMode
  68. // Apply default values
  69. if cfg.FileMode == 0 {
  70. cfg.FileMode = DefaultConfig.FileMode
  71. }
  72. if cfg.DirMode == 0 {
  73. cfg.DirMode = DefaultConfig.DirMode
  74. }
  75. if db.storemode == 2 && db.name == "" {
  76. return db, nil
  77. }
  78. _, err = os.Stat(f)
  79. if err != nil {
  80. // file not exists - create dirs if any
  81. if os.IsNotExist(err) {
  82. if filepath.Dir(f) != "." {
  83. err = os.MkdirAll(filepath.Dir(f), os.FileMode(cfg.DirMode))
  84. if err != nil {
  85. return nil, err
  86. }
  87. }
  88. } else {
  89. return nil, err
  90. }
  91. }
  92. db.fv, err = os.OpenFile(f, os.O_CREATE|os.O_RDWR, os.FileMode(cfg.FileMode))
  93. if err != nil {
  94. return nil, err
  95. }
  96. db.fk, err = os.OpenFile(f+".idx", os.O_CREATE|os.O_RDWR, os.FileMode(cfg.FileMode))
  97. if err != nil {
  98. return nil, err
  99. }
  100. //read keys
  101. buf := new(bytes.Buffer)
  102. b, err := ioutil.ReadAll(db.fk)
  103. if err != nil {
  104. return nil, err
  105. }
  106. buf.Write(b)
  107. var readSeek uint32
  108. for buf.Len() > 0 {
  109. _ = uint8(buf.Next(1)[0]) //format version
  110. t := uint8(buf.Next(1)[0])
  111. seek := binary.BigEndian.Uint32(buf.Next(4))
  112. size := binary.BigEndian.Uint32(buf.Next(4))
  113. _ = buf.Next(4) //time
  114. sizeKey := int(binary.BigEndian.Uint16(buf.Next(2)))
  115. key := buf.Next(sizeKey)
  116. strkey := string(key)
  117. cmd := &Cmd{
  118. Seek: seek,
  119. Size: size,
  120. KeySeek: readSeek,
  121. }
  122. if db.storemode == 2 {
  123. cmd.Val = make([]byte, size)
  124. db.fv.ReadAt(cmd.Val, int64(seek))
  125. }
  126. readSeek += uint32(16 + sizeKey)
  127. switch t {
  128. case 0:
  129. if _, exists := db.vals[strkey]; !exists {
  130. //write new key at keys store
  131. db.appendKey(key)
  132. }
  133. db.vals[strkey] = cmd
  134. case 1:
  135. delete(db.vals, strkey)
  136. db.deleteFromKeys(key)
  137. }
  138. }
  139. if cfg.SyncInterval > 0 {
  140. db.backgroundManager(cfg.SyncInterval)
  141. }
  142. return db, err
  143. }
  144. // backgroundManager runs continuously in the background and performs various
  145. // operations such as syncing to disk.
  146. func (db *Db) backgroundManager(interval int) {
  147. ctx, cancel := context.WithCancel(context.Background())
  148. db.cancelSyncer = cancel
  149. go func() {
  150. for {
  151. select {
  152. case <-ctx.Done():
  153. return
  154. default:
  155. db.Lock()
  156. db.fk.Sync()
  157. db.fv.Sync()
  158. db.Unlock()
  159. time.Sleep(time.Duration(interval) * time.Second)
  160. }
  161. }
  162. }()
  163. }
  164. //appendKey insert key in slice
  165. func (db *Db) appendKey(b []byte) {
  166. //log.Println("append")
  167. db.keys = append(db.keys, b)
  168. return
  169. }
  170. // deleteFromKeys delete key from slice keys
  171. func (db *Db) deleteFromKeys(b []byte) {
  172. found := db.found(b, true)
  173. if found < len(db.keys) {
  174. if bytes.Equal(db.keys[found], b) {
  175. db.keys = append(db.keys[:found], db.keys[found+1:]...)
  176. }
  177. }
  178. }
  179. func (db *Db) sort() {
  180. if !sort.SliceIsSorted(db.keys, db.lessBinary) {
  181. //log.Println("sort")
  182. sort.Slice(db.keys, db.lessBinary)
  183. }
  184. }
  185. func (db *Db) lessBinary(i, j int) bool {
  186. return bytes.Compare(db.keys[i], db.keys[j]) <= 0
  187. }
  188. //found return binary search result with sort order
  189. func (db *Db) found(b []byte, asc bool) int {
  190. db.sort()
  191. //if asc {
  192. return sort.Search(len(db.keys), func(i int) bool {
  193. return bytes.Compare(db.keys[i], b) >= 0
  194. })
  195. //}
  196. //return sort.Search(len(db.keys), func(i int) bool {
  197. // return bytes.Compare(db.keys[i], b) <= 0
  198. //})
  199. }
  200. // KeyToBinary return key in bytes
  201. func KeyToBinary(v interface{}) ([]byte, error) {
  202. var err error
  203. switch v.(type) {
  204. case []byte:
  205. return v.([]byte), nil
  206. case bool, float32, float64, complex64, complex128, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
  207. buf := new(bytes.Buffer)
  208. err = binary.Write(buf, binary.BigEndian, v)
  209. return buf.Bytes(), err
  210. case int:
  211. val := uint64(v.(int))
  212. p := make([]byte, 8)
  213. p[0] = byte(val >> 56)
  214. p[1] = byte(val >> 48)
  215. p[2] = byte(val >> 40)
  216. p[3] = byte(val >> 32)
  217. p[4] = byte(val >> 24)
  218. p[5] = byte(val >> 16)
  219. p[6] = byte(val >> 8)
  220. p[7] = byte(val)
  221. return p, err
  222. case string:
  223. return []byte(v.(string)), nil
  224. default:
  225. buf := new(bytes.Buffer)
  226. err = gob.NewEncoder(buf).Encode(v)
  227. return buf.Bytes(), err
  228. }
  229. }
  230. // ValToBinary return value in bytes
  231. func ValToBinary(v interface{}) ([]byte, error) {
  232. var err error
  233. switch v.(type) {
  234. case []byte:
  235. return v.([]byte), nil
  236. default:
  237. buf := new(bytes.Buffer)
  238. err = gob.NewEncoder(buf).Encode(v)
  239. if err != nil {
  240. return nil, err
  241. }
  242. return buf.Bytes(), err
  243. }
  244. }
  245. func writeKeyVal(fk, fv *os.File, readKey, writeVal []byte, exists bool, oldCmd *Cmd) (cmd *Cmd, err error) {
  246. var seek, newSeek int64
  247. cmd = &Cmd{Size: uint32(len(writeVal))}
  248. if exists {
  249. // key exists
  250. cmd.Seek = oldCmd.Seek
  251. cmd.KeySeek = oldCmd.KeySeek
  252. if oldCmd.Size >= uint32(len(writeVal)) {
  253. //write at old seek new value
  254. _, _, err = writeAtPos(fv, writeVal, int64(oldCmd.Seek))
  255. } else {
  256. //write at new seek (at the end of file)
  257. seek, _, err = writeAtPos(fv, writeVal, int64(-1))
  258. cmd.Seek = uint32(seek)
  259. }
  260. if err == nil {
  261. // if no error - store key at KeySeek
  262. newSeek, err = writeKey(fk, 0, cmd.Seek, cmd.Size, []byte(readKey), int64(cmd.KeySeek))
  263. cmd.KeySeek = uint32(newSeek)
  264. }
  265. } else {
  266. // new key
  267. // write value at the end of file
  268. seek, _, err = writeAtPos(fv, writeVal, int64(-1))
  269. cmd.Seek = uint32(seek)
  270. if err == nil {
  271. newSeek, err = writeKey(fk, 0, cmd.Seek, cmd.Size, []byte(readKey), -1)
  272. cmd.KeySeek = uint32(newSeek)
  273. }
  274. }
  275. return cmd, err
  276. }
  277. // if pos<0 store at the end of file
  278. func writeAtPos(f *os.File, b []byte, pos int64) (seek int64, n int, err error) {
  279. seek = pos
  280. if pos < 0 {
  281. seek, err = f.Seek(0, 2)
  282. if err != nil {
  283. return seek, 0, err
  284. }
  285. }
  286. n, err = f.WriteAt(b, seek)
  287. if err != nil {
  288. return seek, n, err
  289. }
  290. return seek, n, err
  291. }
  292. // writeKey create buffer and store key with val address and size
  293. func writeKey(fk *os.File, t uint8, seek, size uint32, key []byte, keySeek int64) (newSeek int64, err error) {
  294. //get buf from pool
  295. buf := new(bytes.Buffer)
  296. buf.Reset()
  297. buf.Grow(16 + len(key))
  298. //encode
  299. binary.Write(buf, binary.BigEndian, uint8(0)) //1byte version
  300. binary.Write(buf, binary.BigEndian, t) //1byte command code(0-set,1-delete)
  301. binary.Write(buf, binary.BigEndian, seek) //4byte seek
  302. binary.Write(buf, binary.BigEndian, size) //4byte size
  303. binary.Write(buf, binary.BigEndian, uint32(time.Now().Unix())) //4byte timestamp
  304. binary.Write(buf, binary.BigEndian, uint16(len(key))) //2byte key size
  305. buf.Write(key) //key
  306. if keySeek < 0 {
  307. newSeek, _, err = writeAtPos(fk, buf.Bytes(), int64(-1))
  308. } else {
  309. newSeek, _, err = writeAtPos(fk, buf.Bytes(), int64(keySeek))
  310. }
  311. return newSeek, err
  312. }
  313. // findKey return index of first key in ascending mode
  314. // findKey return index of last key in descending mode
  315. // findKey return 0 or len-1 in case of nil key
  316. func (db *Db) findKey(key interface{}, asc bool) (int, error) {
  317. if key == nil {
  318. db.sort()
  319. if asc {
  320. return 0, ErrKeyNotFound
  321. }
  322. return len(db.keys) - 1, ErrKeyNotFound
  323. }
  324. k, err := KeyToBinary(key)
  325. if err != nil {
  326. return -1, err
  327. }
  328. found := db.found(k, asc)
  329. //log.Println("found", found)
  330. // check found
  331. if found >= len(db.keys) {
  332. return -1, ErrKeyNotFound
  333. }
  334. if !bytes.Equal(db.keys[found], k) {
  335. return -1, ErrKeyNotFound
  336. }
  337. return found, nil
  338. }
  339. // startFrom return is a start from b in binary
  340. func startFrom(a, b []byte) bool {
  341. if a == nil || b == nil {
  342. return false
  343. }
  344. if len(a) < len(b) {
  345. return false
  346. }
  347. return bytes.Compare(a[:len(b)], b) == 0
  348. }
  349. func (db *Db) foundPref(b []byte, asc bool) int {
  350. db.sort()
  351. if asc {
  352. return sort.Search(len(db.keys), func(i int) bool {
  353. return bytes.Compare(db.keys[i], b) >= 0
  354. })
  355. }
  356. var j int
  357. for j = len(db.keys) - 1; j >= 0; j-- {
  358. if startFrom(db.keys[j], b) {
  359. break
  360. }
  361. }
  362. return j
  363. }
  364. func checkInterval(find, limit, offset, excludeFrom, len int, asc bool) (int, int) {
  365. end := 0
  366. start := find
  367. if asc {
  368. start += (offset + excludeFrom)
  369. if limit == 0 {
  370. end = len - excludeFrom
  371. } else {
  372. end = (start + limit - 1)
  373. }
  374. } else {
  375. start -= (offset + excludeFrom)
  376. if limit == 0 {
  377. end = 0
  378. } else {
  379. end = start - limit + 1
  380. }
  381. }
  382. if end < 0 {
  383. end = 0
  384. }
  385. if end >= len {
  386. end = len - 1
  387. }
  388. return start, end
  389. }