stress_container.py

#!/usr/bin/env python
# Copyright (c) 2021 Nvidia
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
from argparse import ArgumentParser
from os.path import join, exists
from multiprocessing import Pool

from swift.common.ring import Ring
from swift.common.direct_client import direct_put_container_object
from pprint import pprint

node = part = account = container = conn_timeout = None

def put_obj_into_container(num):
    global part, account, container, node, conn_timeout
    sys.stdout.write('.')
    sys.stdout.flush()
    try:
        direct_put_container_object(node, part, account, container, 'obj%d' % num,
                headers={"X-Size": 0, 'X-Content-Type': 'application/stress', 'x-etag': 'something'},
                conn_timeout=int(conn_timeout))
    except Exception as ex:
        print(str(ex))


def main(args):
    ring = Ring(ring_path)
    global part, node, account, container, conn_timeout
    part, nodes = ring.get_nodes(args.account, args.container)
    node = nodes[0]
    account, container, conn_timeout = args.account, args.container, args.conn_timeout
    pprint(node)
    print('part = %d' % part)
    pool = Pool(int(args.num_connections))

    pool.map(put_obj_into_container, range(int(args.total)))
    pool.close()


if __name__ == '__main__':
    usage = '''
    Probes and gethers all shard data.
    Usage: %prog <account> <container>
    '''
    parser = ArgumentParser(usage)
    parser.add_argument('-d', '--swift-dir', default='/etc/swift',
                        dest='swift_dir', help='Path to swift directory')
    parser.add_argument('-n', '--number', default=5, dest='num_connections',
                        help='Number of concurrent connections')
    parser.add_argument('-t', '--total', default=1000, dest='total',
                        help='Total number of puts')
    parser.add_argument('-T', '--connect-timeout', default=5, dest='conn_timeout',
                        help='Direct client connection timeout')
    parser.add_argument('account', type=str, help='account')
    parser.add_argument('container', type=str, help='container')

    args = parser.parse_args()
    ring_path = join(args.swift_dir, 'container.ring.gz')
    if not exists(ring_path):
        print("Container ring '%s' doesn't exist" % ring_path)
        sys.exit(1)

    main(args)