small-obj-load-test.py

from minio import Minio
from minio.error import ResponseError
from sys import stdout
from itertools import cycle

import string
import sys
import datetime
import random
import argparse

oneMB = 1024*1024
OBJECT_SIZE = 65*oneMB
BUCKET = "test"
HOST=""
AKEY="minio"
SKEY="minio123"

def repeat_with_length(s, length=0):
    n = len(s)
    if n == 0 or length < 0:
        raise ValueError("Input string cannot be empty and length has to >= 0.")
    q = length // n
    r = length % n
    return b''.join([s * q, s[:r]])

# creates a file-like object with content being the repetition of a
# given string.
def repeater_file_with_length(s, length=0):
    class MyReader(object):
        consumed = 0
        st_idx = 0
        length = 0
        s = b''

        def __init__(self, s, length):
            self.length = length
            self.s = s

        def read(self, size=-1):
            if self.consumed < self.length:
                if self.st_idx == 0:
                    str_to_repeat = self.s
                else:
                    str_to_repeat = self.s[self.st_idx:] + self.s[:self.st_idx]

                if size == -1:
                    res = repeat_with_length(
                        str_to_repeat,
                        self.length - self.consumed
                    )
                    self.length = self.consumed
                    return res
                if self.consumed + size > self.length:
                    size = self.length - self.consumed

                self.st_idx = (self.st_idx + size) % len(self.s)
                self.consumed += size
                return repeat_with_length(str_to_repeat, size)
            return None

    return MyReader(s, length)

base_content = b'123456789\n'

def upload_it(num):
    fname = "my-boring-object-{}".format(num)
    h = HOST
    path = "{}:{}/{}".format(h, BUCKET, fname)
    mc = Minio(
        HOST,
        access_key=AKEY,
        secret_key=SKEY,
        secure=False
    )
    try:
        l = OBJECT_SIZE
        start_time = datetime.datetime.utcnow()
        print("DEBUG:", BUCKET, fname)
        mc.put_object(BUCKET, fname, repeater_file_with_length(base_content, l), l)
        duration = datetime.datetime.utcnow() - start_time
        r = "{}: upload success - took {}".format(path, duration)
    except ResponseError as err:
        r = "{}: upload failure - got error: {}".format(path, err)
    return r

def upload_objects(count=1):
    for i in range(count):
        res = upload_it(i)
        print(res)

def verify_chunk(data, offset):
    n = len(data)
    i = offset % len(base_content)
    if i == 0:
        bstr = base_content
    else:
        bstr = base_content[i:] + base_content[:i]

    for i, (d, refval) in enumerate(zip(data, cycle(bstr))):
        if d != refval:
            return offset + i
    return None

def verify_obj_helper(obj):
    return get_chunk_and_verify(obj.object_name)

chunk_size = 16*oneMB
def get_chunk_and_verify(fname):
    print("Found file {}. Verifying it.".format(fname))
    h = HOST
    path = "{}:{}/{}".format(h, BUCKET, fname)
    mc = Minio(
        h,
        access_key=AKEY,
        secret_key=SKEY,
        secure=False
    )
    try:
        offset = 0
        chunk_size = 16*oneMB
        while offset < OBJECT_SIZE:
            res = mc.get_partial_object(BUCKET, fname, offset, chunk_size)
            data = res.read()
            check_res = verify_chunk(data, offset)
            if check_res is not None:
                return "{}: Verification failure at byte {}".format(path, check_res)

            offset += chunk_size
        return None
    except ResponseErr as err:
        return "{}: Response error: {}".format(path, err)


def verify_driver():
    global HOST
    mc = Minio(
        HOST,
        access_key=AKEY,
        secret_key=SKEY,
        secure=False
    )

    objects = mc.list_objects(BUCKET, "my-boring-object-", False)
    for obj in objects:
        try:
            res = verify_obj_helper(obj)
            if res is not None:
                print("Failed: {}".format(r))
                
        except ResponseErr as err:
            return "Response error: {}".format(err)



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--count", type=int, nargs=1,
		        help="number of objects to be uploaded")
    parser.add_argument("--verify", help="enable object verification",
			action="store_true")
    parser.add_argument("--host", nargs=1, help="minio host to connect to")
    args = parser.parse_args()
    HOST=args.host[0]
    if args.count is not None:
        count = args.count[0]
        print("count = ", count)
        upload_objects(count)
        
    if args.verify:
        verify_driver()