-
Notifications
You must be signed in to change notification settings - Fork 27
/
mc_bin_client.py
399 lines (316 loc) · 14.8 KB
/
mc_bin_client.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
#!/usr/bin/env python
"""
Binary memcached test client.
Copyright (c) 2007 Dustin Sallings <dustin@spy.net>
"""
import sys
import time
import hmac
import socket
import random
import struct
import exceptions
from memcacheConstants import REQ_MAGIC_BYTE, RES_MAGIC_BYTE
from memcacheConstants import REQ_PKT_FMT, RES_PKT_FMT, MIN_RECV_PACKET
from memcacheConstants import SET_PKT_FMT, DEL_PKT_FMT, INCRDECR_RES_FMT
from memcacheConstants import TOUCH_PKT_FMT, GAT_PKT_FMT, GETL_PKT_FMT
import memcacheConstants
class MemcachedError(exceptions.Exception):
"""Error raised when a command fails."""
def __init__(self, status, msg):
supermsg='Memcached error #' + `status`
if msg: supermsg += ": " + msg
exceptions.Exception.__init__(self, supermsg)
self.status=status
self.msg=msg
def __repr__(self):
return "<MemcachedError #%d ``%s''>" % (self.status, self.msg)
class MemcachedClient(object):
"""Simple memcached client."""
vbucketId = 0
def __init__(self, host='127.0.0.1', port=11211, family=socket.AF_INET):
self.host = host
self.port = port
self.s=socket.socket(family, socket.SOCK_STREAM)
if hasattr(socket, 'AF_UNIX') and family == socket.AF_UNIX:
self.s.connect_ex(host)
else:
self.s.connect_ex((host, port))
self.r=random.Random()
def close(self):
self.s.close()
def __del__(self):
self.close()
def _sendCmd(self, cmd, key, val, opaque, extraHeader='', cas=0):
self._sendMsg(cmd, key, val, opaque, extraHeader=extraHeader, cas=cas,
vbucketId=self.vbucketId)
def _sendMsg(self, cmd, key, val, opaque, extraHeader='', cas=0,
dtype=0, vbucketId=0,
fmt=REQ_PKT_FMT, magic=REQ_MAGIC_BYTE):
msg=struct.pack(fmt, magic,
cmd, len(key), len(extraHeader), dtype, vbucketId,
len(key) + len(extraHeader) + len(val), opaque, cas)
self.s.send(msg + extraHeader + key + val)
def _recvMsg(self):
response = ""
while len(response) < MIN_RECV_PACKET:
data = self.s.recv(MIN_RECV_PACKET - len(response))
if data == '':
raise exceptions.EOFError("Got empty data (remote died?).")
response += data
assert len(response) == MIN_RECV_PACKET
magic, cmd, keylen, extralen, dtype, errcode, remaining, opaque, cas=\
struct.unpack(RES_PKT_FMT, response)
rv = ""
while remaining > 0:
data = self.s.recv(remaining)
if data == '':
raise exceptions.EOFError("Got empty data (remote died?).")
rv += data
remaining -= len(data)
assert (magic in (RES_MAGIC_BYTE, REQ_MAGIC_BYTE)), "Got magic: %d" % magic
return cmd, errcode, opaque, cas, keylen, extralen, rv
def _handleKeyedResponse(self, myopaque):
cmd, errcode, opaque, cas, keylen, extralen, rv = self._recvMsg()
assert myopaque is None or opaque == myopaque, \
"expected opaque %x, got %x" % (myopaque, opaque)
if errcode != 0:
raise MemcachedError(errcode, rv)
return cmd, opaque, cas, keylen, extralen, rv
def _handleSingleResponse(self, myopaque):
cmd, opaque, cas, keylen, extralen, data = self._handleKeyedResponse(myopaque)
return opaque, cas, data
def _doCmd(self, cmd, key, val, extraHeader='', cas=0):
"""Send a command and await its response."""
opaque=self.r.randint(0, 2**32)
self._sendCmd(cmd, key, val, opaque, extraHeader, cas)
return self._handleSingleResponse(opaque)
def _mutate(self, cmd, key, exp, flags, cas, val):
return self._doCmd(cmd, key, val, struct.pack(SET_PKT_FMT, flags, exp),
cas)
def _cat(self, cmd, key, cas, val):
return self._doCmd(cmd, key, val, '', cas)
def append(self, key, value, cas=0):
return self._cat(memcacheConstants.CMD_APPEND, key, cas, value)
def prepend(self, key, value, cas=0):
return self._cat(memcacheConstants.CMD_PREPEND, key, cas, value)
def __incrdecr(self, cmd, key, amt, init, exp):
something, cas, val=self._doCmd(cmd, key, '',
struct.pack(memcacheConstants.INCRDECR_PKT_FMT, amt, init, exp))
return struct.unpack(INCRDECR_RES_FMT, val)[0], cas
def incr(self, key, amt=1, init=0, exp=0):
"""Increment or create the named counter."""
return self.__incrdecr(memcacheConstants.CMD_INCR, key, amt, init, exp)
def decr(self, key, amt=1, init=0, exp=0):
"""Decrement or create the named counter."""
return self.__incrdecr(memcacheConstants.CMD_DECR, key, amt, init, exp)
def _doMetaCmd(self, cmd, key, value, cas, exp, flags, seqno, remote_cas):
extra = struct.pack('>IIQQ', flags, exp, seqno, remote_cas)
return self._doCmd(cmd, key, value, extra, cas)
def set(self, key, exp, flags, val):
"""Set a value in the memcached server."""
return self._mutate(memcacheConstants.CMD_SET, key, exp, flags, 0, val)
def setWithMeta(self, key, value, exp, flags, seqno, remote_cas):
"""Set a value and its meta data in the memcached server."""
return self._doMetaCmd(memcacheConstants.CMD_SET_WITH_META,
key, value, 0, exp, flags, seqno, remote_cas)
def add(self, key, exp, flags, val):
"""Add a value in the memcached server iff it doesn't already exist."""
return self._mutate(memcacheConstants.CMD_ADD, key, exp, flags, 0, val)
def addWithMeta(self, key, value, exp, flags, seqno, remote_cas):
return self._doMetaCmd(memcacheConstants.CMD_ADD_WITH_META,
key, value, 0, exp, flags, seqno, remote_cas)
def replace(self, key, exp, flags, val):
"""Replace a value in the memcached server iff it already exists."""
return self._mutate(memcacheConstants.CMD_REPLACE, key, exp, flags, 0,
val)
def observe(self, key, vbucket):
"""Observe a key for persistence and replication."""
value = struct.pack('>HH', vbucket, len(key)) + key
opaque, cas, data = self._doCmd(memcacheConstants.CMD_OBSERVE, '', value)
rep_time = (cas & 0xFFFFFFFF)
persist_time = (cas >> 32) & 0xFFFFFFFF
persisted = struct.unpack('>B', data[4+len(key)])[0]
return opaque, rep_time, persist_time, persisted
def __parseGet(self, data, klen=0):
flags=struct.unpack(memcacheConstants.GET_RES_FMT, data[-1][:4])[0]
return flags, data[1], data[-1][4 + klen:]
def get(self, key):
"""Get the value for a given key within the memcached server."""
parts=self._doCmd(memcacheConstants.CMD_GET, key, '')
return self.__parseGet(parts)
def getMeta(self, key):
"""Get the metadata for a given key within the memcached server."""
opaque, cas, data = self._doCmd(memcacheConstants.CMD_GET_META, key, '')
deleted = struct.unpack('>I', data[0:4])[0]
flags = struct.unpack('>I', data[4:8])[0]
exp = struct.unpack('>I', data[8:12])[0]
seqno = struct.unpack('>Q', data[12:20])[0]
return (deleted, flags, exp, seqno, cas)
def getl(self, key, exp=15):
"""Get the value for a given key within the memcached server."""
parts=self._doCmd(memcacheConstants.CMD_GET_LOCKED, key, '',
struct.pack(memcacheConstants.GETL_PKT_FMT, exp))
return self.__parseGet(parts)
def cas(self, key, exp, flags, oldVal, val):
"""CAS in a new value for the given key and comparison value."""
self._mutate(memcacheConstants.CMD_SET, key, exp, flags,
oldVal, val)
def touch(self, key, exp):
"""Touch a key in the memcached server."""
return self._doCmd(memcacheConstants.CMD_TOUCH, key, '',
struct.pack(memcacheConstants.TOUCH_PKT_FMT, exp))
def gat(self, key, exp):
"""Get the value for a given key and touch it within the memcached server."""
parts=self._doCmd(memcacheConstants.CMD_GAT, key, '',
struct.pack(memcacheConstants.GAT_PKT_FMT, exp))
return self.__parseGet(parts)
def getr(self, key):
"""Get the value for a given key in a replica vbucket within the memcached server."""
parts=self._doCmd(memcacheConstants.CMD_GET_REPLICA, key, '')
return self.__parseGet(parts, len(key))
def version(self):
"""Get the value for a given key within the memcached server."""
return self._doCmd(memcacheConstants.CMD_VERSION, '', '')
def verbose(self, level):
"""Set the verbosity level."""
return self._doCmd(memcacheConstants.CMD_VERBOSE, '', '',
extraHeader=struct.pack(">I", level))
def sasl_mechanisms(self):
"""Get the supported SASL methods."""
return set(self._doCmd(memcacheConstants.CMD_SASL_LIST_MECHS,
'', '')[2].split(' '))
def sasl_auth_start(self, mech, data):
"""Start a sasl auth session."""
return self._doCmd(memcacheConstants.CMD_SASL_AUTH, mech, data)
def sasl_auth_plain(self, user, password, foruser=''):
"""Perform plain auth."""
return self.sasl_auth_start('PLAIN', '\0'.join([foruser, user, password]))
def sasl_auth_cram_md5(self, user, password):
"""Start a plan auth session."""
try:
self.sasl_auth_start('CRAM-MD5', '')
except MemcachedError, e:
if e.status != memcacheConstants.ERR_AUTH_CONTINUE:
raise
challenge = e.msg
dig = hmac.HMAC(password, challenge).hexdigest()
return self._doCmd(memcacheConstants.CMD_SASL_STEP, 'CRAM-MD5',
user + ' ' + dig)
def stop_persistence(self):
return self._doCmd(memcacheConstants.CMD_STOP_PERSISTENCE, '', '')
def start_persistence(self):
return self._doCmd(memcacheConstants.CMD_START_PERSISTENCE, '', '')
def set_param(self, key, val, type):
print "setting param:", key, val
type = struct.pack(memcacheConstants.SET_PARAM_FMT, type)
return self._doCmd(memcacheConstants.CMD_SET_PARAM, key, val, type)
def set_vbucket_state(self, vbucket, stateName):
assert isinstance(vbucket, int)
self.vbucketId = vbucket
state = struct.pack(memcacheConstants.VB_SET_PKT_FMT,
memcacheConstants.VB_STATE_NAMES[stateName])
return self._doCmd(memcacheConstants.CMD_SET_VBUCKET_STATE, '', '', state)
def get_vbucket_state(self, vbucket):
assert isinstance(vbucket, int)
self.vbucketId = vbucket
return self._doCmd(memcacheConstants.CMD_GET_VBUCKET_STATE, '', '')
def delete_vbucket(self, vbucket):
assert isinstance(vbucket, int)
self.vbucketId = vbucket
return self._doCmd(memcacheConstants.CMD_DELETE_VBUCKET, '', '')
def evict_key(self, key):
return self._doCmd(memcacheConstants.CMD_EVICT_KEY, key, '')
def getMulti(self, keys):
"""Get values for any available keys in the given iterable.
Returns a dict of matched keys to their values."""
opaqued=dict(enumerate(keys))
terminal=len(opaqued)+10
# Send all of the keys in quiet
for k,v in opaqued.iteritems():
self._sendCmd(memcacheConstants.CMD_GETQ, v, '', k)
self._sendCmd(memcacheConstants.CMD_NOOP, '', '', terminal)
# Handle the response
rv={}
done=False
while not done:
opaque, cas, data=self._handleSingleResponse(None)
if opaque != terminal:
rv[opaqued[opaque]]=self.__parseGet((opaque, cas, data))
else:
done=True
return rv
def setMulti(self, exp, flags, items):
"""Multi-set (using setq).
Give me (key, value) pairs."""
# If this is a dict, convert it to a pair generator
if hasattr(items, 'iteritems'):
items = items.iteritems()
opaqued=dict(enumerate(items))
terminal=len(opaqued)+10
extra=struct.pack(SET_PKT_FMT, flags, exp)
# Send all of the keys in quiet
for opaque,kv in opaqued.iteritems():
self._sendCmd(memcacheConstants.CMD_SETQ, kv[0], kv[1], opaque, extra)
self._sendCmd(memcacheConstants.CMD_NOOP, '', '', terminal)
# Handle the response
failed = []
done=False
while not done:
try:
opaque, cas, data = self._handleSingleResponse(None)
done = opaque == terminal
except MemcachedError, e:
failed.append(e)
return failed
def delMulti(self, items):
"""Multi-delete (using delq).
Give me a collection of keys."""
opaqued = dict(enumerate(items))
terminal = len(opaqued)+10
extra = ''
# Send all of the keys in quiet
for opaque, k in opaqued.iteritems():
self._sendCmd(memcacheConstants.CMD_DELETEQ, k, '', opaque, extra)
self._sendCmd(memcacheConstants.CMD_NOOP, '', '', terminal)
# Handle the response
failed = []
done=False
while not done:
try:
opaque, cas, data = self._handleSingleResponse(None)
done = opaque == terminal
except MemcachedError, e:
failed.append(e)
return failed
def stats(self, sub=''):
"""Get stats."""
opaque=self.r.randint(0, 2**32)
self._sendCmd(memcacheConstants.CMD_STAT, sub, '', opaque)
done = False
rv = {}
while not done:
cmd, opaque, cas, klen, extralen, data = self._handleKeyedResponse(None)
if klen:
rv[data[0:klen]] = data[klen:]
else:
done = True
return rv
def noop(self):
"""Send a noop command."""
return self._doCmd(memcacheConstants.CMD_NOOP, '', '')
def delete(self, key, cas=0):
"""Delete the value for a given key within the memcached server."""
return self._doCmd(memcacheConstants.CMD_DELETE, key, '', '', cas)
def flush(self, timebomb=0):
"""Flush all storage in a memcached instance."""
return self._doCmd(memcacheConstants.CMD_FLUSH, '', '',
struct.pack(memcacheConstants.FLUSH_PKT_FMT, timebomb))
def bucket_select(self, name):
return self._doCmd(memcacheConstants.CMD_SELECT_BUCKET, name, '')
def deregister_tap_client(self, tap_name):
"""Deregister the TAP client with a given name."""
return self._doCmd(memcacheConstants.CMD_DEREGISTER_TAP_CLIENT, tap_name, '', '', 0)
def reset_replication_chain(self):
"""Reset the replication chain."""
return self._doCmd(memcacheConstants.CMD_RESET_REPLICATION_CHAIN, '', '', '', 0)