// Copyright 2019 Lunny Xiao. All rights reserved. // Use of this source code is governed by a MIT-style // license that can be found in the LICENSE file. package levelqueue import ( "bytes" "encoding/binary" "sync" "github.com/syndtr/goleveldb/leveldb" ) // Queue defines a queue struct type Queue struct { db *leveldb.DB highLock sync.Mutex lowLock sync.Mutex low int64 high int64 } // Open opens a queue object or create it if not exist func Open(dataDir string) (*Queue, error) { db, err := leveldb.OpenFile(dataDir, nil) if err != nil { return nil, err } var queue = &Queue{ db: db, } queue.low, err = queue.readID(lowKey) if err == leveldb.ErrNotFound { queue.low = 1 err = db.Put(lowKey, id2bytes(1), nil) } if err != nil { return nil, err } queue.high, err = queue.readID(highKey) if err == leveldb.ErrNotFound { err = db.Put(highKey, id2bytes(0), nil) } if err != nil { return nil, err } return queue, nil } func (queue *Queue) readID(key []byte) (int64, error) { bs, err := queue.db.Get(key, nil) if err != nil { return 0, err } return bytes2id(bs) } var ( lowKey = []byte("low") highKey = []byte("high") ) func (queue *Queue) highincrement() (int64, error) { id := queue.high + 1 queue.high = id err := queue.db.Put(highKey, id2bytes(queue.high), nil) if err != nil { queue.high = queue.high - 1 return 0, err } return id, nil } func (queue *Queue) highdecrement() (int64, error) { queue.high = queue.high - 1 err := queue.db.Put(highKey, id2bytes(queue.high), nil) if err != nil { queue.high = queue.high + 1 return 0, err } return queue.high, nil } func (queue *Queue) lowincrement() (int64, error) { queue.low = queue.low + 1 err := queue.db.Put(lowKey, id2bytes(queue.low), nil) if err != nil { queue.low = queue.low - 1 return 0, err } return queue.low, nil } func (queue *Queue) lowdecrement() (int64, error) { queue.low = queue.low - 1 err := queue.db.Put(lowKey, id2bytes(queue.low), nil) if err != nil { queue.low = queue.low + 1 return 0, err } return queue.low, nil } // Len returns the length of the queue func (queue *Queue) Len() int64 { queue.lowLock.Lock() queue.highLock.Lock() l := queue.high - queue.low + 1 queue.highLock.Unlock() queue.lowLock.Unlock() return l } func id2bytes(id int64) []byte { var buf = make([]byte, 8) binary.PutVarint(buf, id) return buf } func bytes2id(b []byte) (int64, error) { return binary.ReadVarint(bytes.NewReader(b)) } // RPush pushes a data from right of queue func (queue *Queue) RPush(data []byte) error { queue.highLock.Lock() id, err := queue.highincrement() if err != nil { queue.highLock.Unlock() return err } err = queue.db.Put(id2bytes(id), data, nil) queue.highLock.Unlock() return err } // LPush pushes a data from left of queue func (queue *Queue) LPush(data []byte) error { queue.highLock.Lock() id, err := queue.lowdecrement() if err != nil { queue.highLock.Unlock() return err } err = queue.db.Put(id2bytes(id), data, nil) queue.highLock.Unlock() return err } // RPop pop a data from right of queue func (queue *Queue) RPop() ([]byte, error) { queue.highLock.Lock() currentID := queue.high res, err := queue.db.Get(id2bytes(currentID), nil) if err != nil { queue.highLock.Unlock() if err == leveldb.ErrNotFound { return nil, ErrNotFound } return nil, err } _, err = queue.highdecrement() if err != nil { queue.highLock.Unlock() return nil, err } err = queue.db.Delete(id2bytes(currentID), nil) queue.highLock.Unlock() if err != nil { return nil, err } return res, nil } // LPop pop a data from left of queue func (queue *Queue) LPop() ([]byte, error) { queue.lowLock.Lock() currentID := queue.low res, err := queue.db.Get(id2bytes(currentID), nil) if err != nil { queue.lowLock.Unlock() if err == leveldb.ErrNotFound { return nil, ErrNotFound } return nil, err } _, err = queue.lowincrement() if err != nil { return nil, err } err = queue.db.Delete(id2bytes(currentID), nil) queue.lowLock.Unlock() if err != nil { return nil, err } return res, nil } // Close closes the queue func (queue *Queue) Close() error { err := queue.db.Close() queue.db = nil return err }