Statistics
| Branch: | Tag: | Revision:

root / snf-cyclades-app / synnefo / db / pools / __init__.py @ e59cda53

History | View | Annotate | Download (9 kB)

1
from bitarray import bitarray
2
from base64 import b64encode, b64decode
3
import ipaddr
4

    
5
AVAILABLE = True
6
UNAVAILABLE = False
7

    
8

    
9
class PoolManager(object):
10
    """PoolManager for DB PoolTable models.
11

12
    This class implements a persistent Pool mechanism based on rows of
13
    PoolTable objects. Values that are pooled by this class, are mapped to an
14
    index on a bitarray, which is the one that is stored on the DB.
15

16
    The object that will be used in order to initialize this pool, must have
17
    two string attributes (available_map and reserved_map) and the size of the
18
    pool.
19

20
    Subclasses of PoolManager must implement value_to_index and index_to_value
21
    method's in order to denote how the value will be mapped to the index in
22
    the bitarray.
23

24
    Important!!: Updates on a PoolManager object are not reflected to the DB,
25
    until save() method is called.
26

27
    """
28
    def __init__(self, pool_table):
29
        self.pool_table = pool_table
30
        self.pool_size = pool_table.size
31
        if pool_table.available_map:
32
            self.available = _bitarray_from_string(pool_table.available_map)
33
            self.reserved = _bitarray_from_string(pool_table.reserved_map)
34
        else:
35
            self.available = self._create_empty_pool(self.pool_size)
36
            self.reserved = self._create_empty_pool(self.pool_size)
37
            self.add_padding(self.pool_size)
38

    
39
    def _create_empty_pool(self, size):
40
        ba = bitarray(size)
41
        ba.setall(AVAILABLE)
42
        return ba
43

    
44
    def add_padding(self, pool_size):
45
        bits = find_padding(pool_size)
46
        self.available.extend([UNAVAILABLE] * bits)
47
        self.reserved.extend([UNAVAILABLE] * bits)
48

    
49
    def cut_padding(self, pool_size):
50
        bits = find_padding(pool_size)
51
        self.available = self.available[:-bits]
52
        self.reserved = self.reserved[:-bits]
53

    
54
    @property
55
    def pool(self):
56
        return (self.available & self.reserved)
57

    
58
    def get(self, value=None):
59
        """Get a value from the pool."""
60
        if self.empty():
61
            raise EmptyPool
62
        if value is None:
63
            # Get the first available index
64
            index = int(self.pool.index(AVAILABLE))
65
            assert(index < self.pool_size)
66
            self._reserve(index)
67
            return self.index_to_value(index)
68
        else:
69
            if not self.contains(value):
70
                raise InvalidValue("Value %s does not belong to pool." % value)
71
            if self.is_available(value):
72
                self.reserve(value)
73
                return value
74
            else:
75
                raise ValueNotAvailable("Value %s is not available" % value)
76

    
77
    def put(self, value, external=False):
78
        """Return a value to the pool."""
79
        if value is None:
80
            raise ValueError
81
        index = self.value_to_index(value)
82
        self._release(index, external)
83

    
84
    def reserve(self, value, external=False):
85
        """Reserve a value."""
86
        index = self.value_to_index(value)
87
        self._reserve(index, external)
88
        return True
89

    
90
    def save(self, db=True):
91
        """Save changes to the DB."""
92
        self.pool_table.available_map = _bitarray_to_string(self.available)
93
        self.pool_table.reserved_map = _bitarray_to_string(self.reserved)
94
        if db:
95
            self.pool_table.save()
96

    
97
    def empty(self):
98
        """Return True when pool is empty."""
99
        return not self.pool.any()
100

    
101
    def size(self):
102
        """Return the size of the bitarray(original size + padding)."""
103
        return self.pool.length()
104

    
105
    def _reserve(self, index, external=False):
106
        if external:
107
            self.reserved[index] = UNAVAILABLE
108
        else:
109
            self.available[index] = UNAVAILABLE
110

    
111
    def _release(self, index, external=False):
112
        if external:
113
            self.reserved[index] = AVAILABLE
114
        else:
115
            self.available[index] = AVAILABLE
116

    
117
    def contains(self, value):
118
        index = self.value_to_index(value)
119
        return index >= 0 and index < self.pool_size
120

    
121
    def count_available(self):
122
        return self.pool.count(AVAILABLE)
123

    
124
    def count_unavailable(self):
125
        return self.pool_size - self.count_available()
126

    
127
    def count_reserved(self):
128
        return self.reserved[:self.pool_size].count(UNAVAILABLE)
129

    
130
    def count_unreserved(self):
131
        return self.pool_size - self.count_reserved()
132

    
133
    def is_available(self, value, index=False):
134
        if not index:
135
            idx = self.value_to_index(value)
136
        else:
137
            idx = value
138
        return self.pool[idx] == AVAILABLE
139

    
140
    def is_reserved(self, value, index=False):
141
        if not index:
142
            idx = self.value_to_index(value)
143
        else:
144
            idx = value
145
        return self.reserved[idx] == UNAVAILABLE
146

    
147
    def to_01(self):
148
        return self.pool[:self.pool_size].to01()
149

    
150
    def to_map(self):
151
        return self.to_01().replace("0", "X").replace("1", ".")
152

    
153
    def extend(self, bits_num):
154
        assert(bits_num >= 0)
155
        self.resize(bits_num)
156

    
157
    def shrink(self, bits_num):
158
        assert(bits_num >= 0)
159
        size = self.pool_size
160
        tmp = self.available[(size - bits_num): size]
161
        if tmp.count(UNAVAILABLE):
162
            raise Exception("Can not shrink. In use")
163
        self.resize(-bits_num)
164

    
165
    def resize(self, bits_num):
166
        if bits_num == 0:
167
            return
168
        # Cut old padding
169
        self.cut_padding(self.pool_size)
170
        # Do the resize
171
        if bits_num > 0:
172
            self.available.extend([AVAILABLE] * bits_num)
173
            self.reserved.extend([AVAILABLE] * bits_num)
174
        else:
175
            self.available = self.available[:bits_num]
176
            self.reserved = self.reserved[:bits_num]
177
        # Add new padding
178
        self.pool_size = self.pool_size + bits_num
179
        self.add_padding(self.pool_size)
180
        self.pool_table.size = self.pool_size
181

    
182
    def index_to_value(self, index):
183
        raise NotImplementedError
184

    
185
    def value_to_index(self, value):
186
        raise NotImplementedError
187

    
188
    def __repr__(self):
189
        return repr(self.pool_table)
190

    
191

    
192
class EmptyPool(Exception):
193
    pass
194

    
195

    
196
class ValueNotAvailable(Exception):
197
    pass
198

    
199

    
200
class InvalidValue(Exception):
201
    pass
202

    
203

    
204
def find_padding(size):
205
    extra = size % 8
206
    return extra and (8 - extra) or 0
207

    
208

    
209
def bitarray_to_01(bitarray_):
210
    return bitarray_.to01()
211

    
212

    
213
def bitarray_to_map(bitarray_):
214
    return bitarray_to_01(bitarray_).replace("0", "X").replace("1", ".")
215

    
216

    
217
def _bitarray_from_string(bitarray_):
218
    ba = bitarray()
219
    ba.frombytes(b64decode(bitarray_))
220
    return ba
221

    
222

    
223
def _bitarray_to_string(bitarray_):
224
    return b64encode(bitarray_.tobytes())
225

    
226
##
227
## Custom pools
228
##
229

    
230

    
231
class BridgePool(PoolManager):
232
    def index_to_value(self, index):
233
        # Bridge indexes should start from 1
234
        return self.pool_table.base + str(index + 1)
235

    
236
    def value_to_index(self, value):
237
        return int(value.replace(self.pool_table.base, "")) - 1
238

    
239

    
240
class MacPrefixPool(PoolManager):
241
    def __init__(self, pool_table):
242
        do_init = False if pool_table.available_map else True
243
        super(MacPrefixPool, self).__init__(pool_table)
244
        if do_init:
245
            for i in xrange(1, self.pool_size):
246
                if not self.validate_mac(self.index_to_value(i)):
247
                    self._reserve(i, external=True)
248
            # Reserve the first mac-prefix for public-networks
249
            self._reserve(0, external=True)
250

    
251
    def index_to_value(self, index):
252
        """Convert index to mac_prefix"""
253
        base = self.pool_table.base
254
        a = hex(int(base.replace(":", ""), 16) + index).replace("0x", '')
255
        mac_prefix = ":".join([a[x:x + 2] for x in xrange(0, len(a), 2)])
256
        return mac_prefix
257

    
258
    def value_to_index(self, value):
259
        base = self.pool_table.base
260
        return int(value.replace(":", ""), 16) - int(base.replace(":", ""), 16)
261

    
262
    @staticmethod
263
    def validate_mac(value):
264
        hex_ = value.replace(":", "")
265
        bin_ = bin(int(hex_, 16))[2:].zfill(8)
266
        return bin_[6] == '1' and bin_[7] == '0'
267

    
268

    
269
class IPPool(PoolManager):
270
    def __init__(self, pool_table):
271
        subnet = pool_table.subnet
272
        self.net = ipaddr.IPNetwork(subnet.cidr)
273
        if pool_table.base is None:
274
            self.base = pool_table.subnet.cidr
275
        else:
276
            self.base = pool_table.base
277
        if pool_table.offset is None:
278
            self.offset = 0
279
        else:
280
            self.offset = pool_table.offset
281
        super(IPPool, self).__init__(pool_table)
282

    
283
    def value_to_index(self, value):
284
        addr = ipaddr.IPAddress(value)
285
        return int(addr) - int(self.net.network) - int(self.offset)
286

    
287
    def index_to_value(self, index):
288
        return str(self.net[index + int(self.offset)])
289

    
290
    def contains(self, address, index=False):
291
        if index is False:
292
            addr = ipaddr.IPAddress(address)
293
            if addr not in self.net:
294
                return False
295
        return super(IPPool, self).contains(address, index=False)
296

    
297
    def return_start(self):
298
        return str(ipaddr.IPAddress(self.base) + self.offset)
299

    
300
    def return_end(self):
301
        return str(ipaddr.IPAddress(self.base) + self.offset + self.size - 1)