PoolCounter

poolcounter.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
# PoolCounter client
# License: MIT
# Based on https://github.com/wikimedia/operations-debs-python-thumbor-wikimedia/blob/master/wikimedia_thumbor/poolcounter/__init__.py

import socket
import logging
import time

class PoolCounter:
    def __init__(self):
        self.server = 'localhost'
        self.port = 7531
        self.stream = None
        self.logger = logging.getLogger('poolcounter')

    def connect(self):
        self.debug('[PoolCounter] Connecting to: %s %d' % (self.server, self.port))
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        self.stream = s
        return self.stream.connect((self.server, self.port))

    def acq4me(self, key, workers, maxqueue, timeout):
        if not self.stream:
            self.connect()
        try:
            self.debug('[PoolCounter] ACQ4ME %s %d %d %d' % (key, workers, maxqueue, timeout))
            self.stream.send(bytes('ACQ4ME %s %d %d %d\n' % (key, workers, maxqueue, timeout), 'utf-8'))
            data = self.stream.recv(4096).decode('utf-8')
            print(data)
        except socket.error as e:
            self.stream = None
            raise e

        return data == 'LOCKED\n'

    def release(self, key):
        if not self.stream:
            self.connect()

        try:
            self.debug('[PoolCounter] RELEASE')
            self.stream.send(bytes('RELEASE %s\n' % key, 'utf-8'))
            data = self.stream.recv(4096).decode('utf-8')
        except socket.error as e:
            self.stream = None
            raise e

        return data == 'RELEASED\n'

    def close(self):
        if self.stream:
            self.debug('[PoolCounter] Disconnecting')
            self.stream.close()
            self.stream = None
            return True
        else:
            self.debug('[PoolCounter] Already disconnected')
            return False

    def debug(self, message):
        self.logger.debug(message)

poolcounter = PoolCounter()
for i in range(15):
    print(i)
    print(poolcounter.acq4me('wtf', 3, 7, 5))
    #poolcounter.release('wtf')