You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
229 lines
6.5 KiB
229 lines
6.5 KiB
import os
|
|
import glob
|
|
import json
|
|
import struct
|
|
import sqlite3
|
|
from collections import deque
|
|
|
|
|
|
class FifoMemoryQueue(object):
|
|
"""In-memory FIFO queue, API compliant with FifoDiskQueue."""
|
|
|
|
def __init__(self):
|
|
self.q = deque()
|
|
self.push = self.q.append
|
|
|
|
def pop(self):
|
|
q = self.q
|
|
return q.popleft() if q else None
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
def __len__(self):
|
|
return len(self.q)
|
|
|
|
|
|
class LifoMemoryQueue(FifoMemoryQueue):
|
|
"""In-memory LIFO queue, API compliant with LifoDiskQueue."""
|
|
|
|
def pop(self):
|
|
q = self.q
|
|
return q.pop() if q else None
|
|
|
|
|
|
class FifoDiskQueue(object):
|
|
"""Persistent FIFO queue."""
|
|
|
|
szhdr_format = ">L"
|
|
szhdr_size = struct.calcsize(szhdr_format)
|
|
|
|
def __init__(self, path, chunksize=100000):
|
|
self.path = path
|
|
if not os.path.exists(path):
|
|
os.makedirs(path)
|
|
self.info = self._loadinfo(chunksize)
|
|
self.chunksize = self.info['chunksize']
|
|
self.headf = self._openchunk(self.info['head'][0], 'ab+')
|
|
self.tailf = self._openchunk(self.info['tail'][0])
|
|
os.lseek(self.tailf.fileno(), self.info['tail'][2], os.SEEK_SET)
|
|
|
|
def push(self, string):
|
|
if not isinstance(string, bytes):
|
|
raise TypeError('Unsupported type: {}'.format(type(string).__name__))
|
|
hnum, hpos = self.info['head']
|
|
hpos += 1
|
|
szhdr = struct.pack(self.szhdr_format, len(string))
|
|
os.write(self.headf.fileno(), szhdr + string)
|
|
if hpos == self.chunksize:
|
|
hpos = 0
|
|
hnum += 1
|
|
self.headf.close()
|
|
self.headf = self._openchunk(hnum, 'ab+')
|
|
self.info['size'] += 1
|
|
self.info['head'] = [hnum, hpos]
|
|
|
|
def _openchunk(self, number, mode='rb'):
|
|
return open(os.path.join(self.path, 'q%05d' % number), mode)
|
|
|
|
def pop(self):
|
|
tnum, tcnt, toffset = self.info['tail']
|
|
if [tnum, tcnt] >= self.info['head']:
|
|
return
|
|
tfd = self.tailf.fileno()
|
|
szhdr = os.read(tfd, self.szhdr_size)
|
|
if not szhdr:
|
|
return
|
|
size, = struct.unpack(self.szhdr_format, szhdr)
|
|
data = os.read(tfd, size)
|
|
tcnt += 1
|
|
toffset += self.szhdr_size + size
|
|
if tcnt == self.chunksize and tnum <= self.info['head'][0]:
|
|
tcnt = toffset = 0
|
|
tnum += 1
|
|
self.tailf.close()
|
|
os.remove(self.tailf.name)
|
|
self.tailf = self._openchunk(tnum)
|
|
self.info['size'] -= 1
|
|
self.info['tail'] = [tnum, tcnt, toffset]
|
|
return data
|
|
|
|
def close(self):
|
|
self.headf.close()
|
|
self.tailf.close()
|
|
self._saveinfo(self.info)
|
|
if len(self) == 0:
|
|
self._cleanup()
|
|
|
|
def __len__(self):
|
|
return self.info['size']
|
|
|
|
def _loadinfo(self, chunksize):
|
|
infopath = self._infopath()
|
|
if os.path.exists(infopath):
|
|
with open(infopath) as f:
|
|
info = json.load(f)
|
|
else:
|
|
info = {
|
|
'chunksize': chunksize,
|
|
'size': 0,
|
|
'tail': [0, 0, 0],
|
|
'head': [0, 0],
|
|
}
|
|
return info
|
|
|
|
def _saveinfo(self, info):
|
|
with open(self._infopath(), 'w') as f:
|
|
json.dump(info, f)
|
|
|
|
def _infopath(self):
|
|
return os.path.join(self.path, 'info.json')
|
|
|
|
def _cleanup(self):
|
|
for x in glob.glob(os.path.join(self.path, 'q*')):
|
|
os.remove(x)
|
|
os.remove(os.path.join(self.path, 'info.json'))
|
|
if not os.listdir(self.path):
|
|
os.rmdir(self.path)
|
|
|
|
|
|
|
|
class LifoDiskQueue(object):
|
|
"""Persistent LIFO queue."""
|
|
|
|
SIZE_FORMAT = ">L"
|
|
SIZE_SIZE = struct.calcsize(SIZE_FORMAT)
|
|
|
|
def __init__(self, path):
|
|
self.path = path
|
|
if os.path.exists(path):
|
|
self.f = open(path, 'rb+')
|
|
qsize = self.f.read(self.SIZE_SIZE)
|
|
self.size, = struct.unpack(self.SIZE_FORMAT, qsize)
|
|
self.f.seek(0, os.SEEK_END)
|
|
else:
|
|
self.f = open(path, 'wb+')
|
|
self.f.write(struct.pack(self.SIZE_FORMAT, 0))
|
|
self.size = 0
|
|
|
|
def push(self, string):
|
|
if not isinstance(string, bytes):
|
|
raise TypeError('Unsupported type: {}'.format(type(string).__name__))
|
|
self.f.write(string)
|
|
ssize = struct.pack(self.SIZE_FORMAT, len(string))
|
|
self.f.write(ssize)
|
|
self.size += 1
|
|
|
|
def pop(self):
|
|
if not self.size:
|
|
return
|
|
self.f.seek(-self.SIZE_SIZE, os.SEEK_END)
|
|
size, = struct.unpack(self.SIZE_FORMAT, self.f.read())
|
|
self.f.seek(-size-self.SIZE_SIZE, os.SEEK_END)
|
|
data = self.f.read(size)
|
|
self.f.seek(-size, os.SEEK_CUR)
|
|
self.f.truncate()
|
|
self.size -= 1
|
|
return data
|
|
|
|
def close(self):
|
|
if self.size:
|
|
self.f.seek(0)
|
|
self.f.write(struct.pack(self.SIZE_FORMAT, self.size))
|
|
self.f.close()
|
|
if not self.size:
|
|
os.remove(self.path)
|
|
|
|
def __len__(self):
|
|
return self.size
|
|
|
|
|
|
class FifoSQLiteQueue(object):
|
|
|
|
_sql_create = (
|
|
'CREATE TABLE IF NOT EXISTS queue '
|
|
'(id INTEGER PRIMARY KEY AUTOINCREMENT, item BLOB)'
|
|
)
|
|
_sql_size = 'SELECT COUNT(*) FROM queue'
|
|
_sql_push = 'INSERT INTO queue (item) VALUES (?)'
|
|
_sql_pop = 'SELECT id, item FROM queue ORDER BY id LIMIT 1'
|
|
_sql_del = 'DELETE FROM queue WHERE id = ?'
|
|
|
|
def __init__(self, path):
|
|
self._path = os.path.abspath(path)
|
|
self._db = sqlite3.Connection(self._path, timeout=60)
|
|
self._db.text_factory = bytes
|
|
with self._db as conn:
|
|
conn.execute(self._sql_create)
|
|
|
|
def push(self, item):
|
|
if not isinstance(item, bytes):
|
|
raise TypeError('Unsupported type: {}'.format(type(item).__name__))
|
|
|
|
with self._db as conn:
|
|
conn.execute(self._sql_push, (item,))
|
|
|
|
def pop(self):
|
|
with self._db as conn:
|
|
for id_, item in conn.execute(self._sql_pop):
|
|
conn.execute(self._sql_del, (id_,))
|
|
return item
|
|
|
|
def close(self):
|
|
size = len(self)
|
|
self._db.close()
|
|
if not size:
|
|
os.remove(self._path)
|
|
|
|
def __len__(self):
|
|
with self._db as conn:
|
|
return next(conn.execute(self._sql_size))[0]
|
|
|
|
|
|
class LifoSQLiteQueue(FifoSQLiteQueue):
|
|
|
|
_sql_pop = 'SELECT id, item FROM queue ORDER BY id DESC LIMIT 1'
|
|
|
|
|
|
#FifoDiskQueue = FifoSQLiteQueue # noqa
|
|
#LifoDiskQueue = LifoSQLiteQueue # noqa
|