memcached_reproduce.py

# coding: utf-8

import socket

switch = 'append'  # you can modify this

_REQ_TMPL = "%s test_profile_key_%%s 0 0 4\r\n2333\r\n"
if switch == 'set':
    REQ_TMPL = _REQ_TMPL % 'set'
    RES = 'STORED\r\n'
elif switch == 'add':
    REQ_TMPL = _REQ_TMPL % 'set'
    RES = 'STORED\r\n'
elif switch == 'replace':
    REQ_TMPL = _REQ_TMPL % 'replace'
    RES = 'NOT_STORED\r\n'
elif switch == 'append':
    REQ_TMPL = _REQ_TMPL % 'append'
    RES = 'NOT_STORED\r\n'
elif switch == 'prepend':
    REQ_TMPL = _REQ_TMPL % 'prepend'
    RES = 'NOT_STORED\r\n'
else:
    raise NotImplementedError()


def hello():
    sock = socket.socket()
    sock.connect(("127.0.0.1", 11211))
    sock.send('version\r\n')
    rv = ''
    while not rv:
        rv += sock.recv(1024)
    print rv
    sock.send('flush_all\r\n')
    rv = sock.recv(1024)
    assert rv == 'OK\r\n'
    sock.send(REQ_TMPL % 'inf')
    rv = sock.recv(1024)
    assert rv == RES, "%s != %s" % (rv, RES)
    sock.close()


def reproduce(n):
    sock = socket.socket()
    sock.connect(("127.0.0.1", 11211))
    out_list = [REQ_TMPL % i for i in range(n)]
    out_buffer = ''.join(out_list)
    n_to_send = len(out_buffer)
    n_sent = 0
    print 'begin send'
    while n_sent < n_to_send:
        n_sent += sock.send(out_buffer[n_sent:])
        # print 'send [%d/%d]' % (n_sent, n_to_send)
    print 'end send'

    n_to_recv = len(RES) * n
    in_buffer = ''
    print 'begin recv'
    while n_to_recv > 0:
        tmp_buf = sock.recv(n_to_recv)
        n_to_recv -= len(tmp_buf)
        in_buffer += tmp_buf
        # print 'recv [%d/%d]' % (len(in_buffer), len(RES) * n)
    print 'end recv'
    assert in_buffer == RES * n
    sock.close()


def main():
    print 'switch: %s' % switch
    hello()
    reproduce(70000)


if __name__ == '__main__':
    main()