Skip to content

Instantly share code, notes, and snippets.

@dustin
Created August 10, 2010 07:56
Show Gist options
  • Save dustin/516878 to your computer and use it in GitHub Desktop.
Save dustin/516878 to your computer and use it in GitHub Desktop.
transfer simulator
#!/usr/bin/env python
"""
tap protocol client.
Copyright (c) 2010 Dustin Sallings <dustin@spy.net>
"""
import sys
import time
import struct
import signal
import asyncore
import mc_bin_client
import tap
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
import memcacheConstants
VBUCKETS=list(range(16))
def setupVbuckets(tdesc):
signal.alarm(5)
mc = mc_bin_client.MemcachedClient(tdesc.host, tdesc.port)
for i in VBUCKETS:
print "Setting vbucket %d state to active." % i
mc.set_vbucket_state(i, 'active')
signal.alarm(0)
def alarm_handler(*stuff):
print "ALRM rang"
signal.pause()
class Handler(object):
def __init__(self, myid):
self.seen = 0
self.vbsets = 0
self.myid = myid
def __call__(self, identifier, cmd, extra, key, vb, val, cas):
signal.alarm(10)
assert vb == self.myid
if cmd in [memcacheConstants.CMD_TAP_VBUCKET_SET]:
state = struct.unpack("!I", val)[0]
print "vb set (%d):\t%s" % (self.myid, ["0", "active", "2", "pending"][state])
self.vbsets += 1
else:
self.seen += 1
if __name__ == '__main__':
signal.signal(signal.SIGALRM, alarm_handler)
descriptors = [tap.TapDescriptor(a) for a in sys.argv[1:]]
for i in range(1000):
for tdesc in descriptors:
setupVbuckets(tdesc)
callbacks = []
for v in VBUCKETS:
cb = Handler(v)
callbacks.append(cb)
opts = {memcacheConstants.TAP_FLAG_TAKEOVER_VBUCKETS: '',
memcacheConstants.TAP_FLAG_LIST_VBUCKETS: [v]}
tap.TapClient(descriptors, cb, opts=opts)
asyncore.loop()
print "Read %d items" % sum(cb.seen for cb in callbacks)
assert sum(cb.vbsets for cb in callbacks) == len(VBUCKETS) * 2
signal.alarm(0)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment