1#!/usr/bin/env python 2# Copyright 2015 Amazon.com, Inc. or its affiliates. All Rights Reserved. 3# 4# Licensed under the Apache License, Version 2.0 (the "License"). You 5# may not use this file except in compliance with the License. A copy of 6# the License is located at 7# 8# http://aws.amazon.com/apache2.0/ 9# 10# or in the "license" file accompanying this file. This file is 11# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF 12# ANY KIND, either express or implied. See the License for the specific 13# language governing permissions and limitations under the License. 14"""Driver for client scripts. 15 16How it Works 17============ 18 19This script is part of the infrastructure used for resource leak tests. 20At a high level, we're trying to verify that specific types of operations 21don't leak memory. For example: 22 23* Creating 100 clients in a loop doesn't leak memory within reason 24* Making 100 API calls doesn't leak memory 25* Streaming uploads/downloads to/from S3 have O(1) memory usage. 26 27In order to do this, these tests are written to utilize two separate 28processes: a driver and a runner. The driver sends commands to the 29runner which performs the actual commands. While the runner is 30running commands the driver can record various attributes about 31the runner process. So far this is just memory usage. 32 33There's a BaseClientDriverTest test that implements the driver part 34of the tests. These should read like normal functional/integration tests. 35 36This script (cmd-runner) implements the runner. It listens for commands 37from the driver and runs the commands. 38 39The driver and runner communicate via a basic text protocol. Commands 40are line separated with arguments separated by spaces. Commands 41are sent to the runner via stdin. 42 43On success, the runner response through stdout with "OK". 44 45Here's an example:: 46 47 +--------+ +-------+ 48 | Driver | |Runner | 49 +--------+ +-------+ 50 51 create_client s3 52 --------------------------------------------> 53 OK 54 <------------------------------------------- 55 56 make_aws_request 100 ec2 describe_instances 57 --------------------------------------------> 58 OK 59 <------------------------------------------- 60 61 stream_s3_upload bucket key /tmp/foo.txt 62 --------------------------------------------> 63 OK 64 <------------------------------------------- 65 66 exit 67 --------------------------------------------> 68 OK 69 <------------------------------------------- 70 71""" 72import botocore.session 73import sys 74 75 76class CommandRunner(object): 77 DEFAULT_REGION = 'us-west-2' 78 79 def __init__(self): 80 self._session = botocore.session.get_session() 81 self._clients = [] 82 self._waiters = [] 83 self._paginators = [] 84 85 def run(self): 86 while True: 87 line = sys.stdin.readline() 88 parts = line.strip().split() 89 if not parts: 90 break 91 elif parts[0] == 'exit': 92 sys.stdout.write('OK\n') 93 sys.stdout.flush() 94 break 95 else: 96 getattr(self, '_do_%s' % parts[0])(parts[1:]) 97 sys.stdout.write('OK\n') 98 sys.stdout.flush() 99 100 def _do_create_client(self, args): 101 # The args provided by the user map directly to the args 102 # passed to create_client. At the least you need 103 # to provide the service name. 104 self._clients.append(self._session.create_client(*args)) 105 106 def _do_create_multiple_clients(self, args): 107 # Args: 108 # * num_clients - Number of clients to create in a loop 109 # * client_args - Variable number of args to pass to create_client 110 num_clients = args[0] 111 client_args = args[1:] 112 for _ in range(int(num_clients)): 113 self._clients.append(self._session.create_client(*client_args)) 114 115 def _do_free_clients(self, args): 116 # Frees the memory associated with all clients. 117 self._clients = [] 118 119 def _do_create_waiter(self, args): 120 # Args: 121 # [0] client name used in its instantiation 122 # [1] waiter name used in its instantiation 123 client = self._create_client(args[0]) 124 self._waiters.append(client.get_waiter(args[1])) 125 126 def _do_create_multiple_waiters(self, args): 127 # Args: 128 # [0] num_waiters - Number of clients to create in a loop 129 # [1] client name used in its instantiation 130 # [2] waiter name used in its instantiation 131 num_waiters = args[0] 132 client = self._create_client(args[1]) 133 for _ in range(int(num_waiters)): 134 self._waiters.append(client.get_waiter(args[2])) 135 136 def _do_free_waiters(self, args): 137 # Frees the memory associated with all of the waiters. 138 self._waiters = [] 139 140 def _do_create_paginator(self, args): 141 # Args: 142 # [0] client name used in its instantiation 143 # [1] paginator name used in its instantiation 144 client = self._create_client(args[0]) 145 self._paginators.append(client.get_paginator(args[1])) 146 147 def _do_create_multiple_paginators(self, args): 148 # Args: 149 # [0] num_paginators - Number of paginators to create in a loop 150 # [1] client name used in its instantiation 151 # [2] paginator name used in its instantiation 152 num_paginators = args[0] 153 client = self._create_client(args[1]) 154 for _ in range(int(num_paginators)): 155 self._paginators.append(client.get_paginator(args[2])) 156 157 def _do_free_paginators(self, args): 158 # Frees the memory associated with all of the waiters. 159 self._paginators = [] 160 161 def _do_make_aws_request(self, args): 162 # Create a client and make a number of AWS requests. 163 # Args: 164 # * num_requests - The number of requests to create in a loop 165 # * service_name - The name of the AWS service 166 # * operation_name - The name of the service, snake_cased 167 # * oepration_args - Variable args, kwargs to pass to the API call, 168 # in the format Kwarg1:Val1 Kwarg2:Val2 169 num_requests = int(args[0]) 170 service_name = args[1] 171 operation_name = args[2] 172 kwargs = dict([v.split(':', 1) for v in args[3:]]) 173 client = self._create_client(service_name) 174 method = getattr(client, operation_name) 175 for _ in range(num_requests): 176 method(**kwargs) 177 178 def _do_stream_s3_upload(self, args): 179 # Stream an upload to S3 from a local file. 180 # This does *not* create an S3 bucket. You need to create this 181 # before running this command. You will also need to create 182 # the local file to upload before calling this command. 183 # Args: 184 # * bucket - The name of the S3 bucket 185 # * key - The name of the S3 key 186 # * filename - The name of the local filename to upload 187 bucket, key, local_filename = args 188 client = self._create_client('s3') 189 with open(local_filename, 'rb') as f: 190 client.put_object(Bucket=bucket, Key=key, Body=f) 191 192 def _do_stream_s3_download(self, args): 193 # Stream a download to S3 from a local file. 194 # Before calling this command you'll need to create the S3 bucket 195 # as well as the S3 object. Also, the directory where the 196 # file will be downloaded must already exist. 197 # Args: 198 # * bucket - The name of the S3 bucket 199 # * key - The name of the S3 key 200 # * filename - The local filename where the object will be downloaded 201 bucket, key, local_filename = args 202 client = self._create_client('s3') 203 response = client.get_object(Bucket=bucket, Key=key) 204 body_stream = response['Body'] 205 with open(local_filename, 'wb') as f: 206 for chunk in iter(lambda: body_stream.read(64 * 1024), b''): 207 f.write(chunk) 208 209 def _create_client(self, service_name): 210 # Create a client using the provided service name. 211 # It will also inject a region name of self.DEFAULT_REGION. 212 return self._session.create_client(service_name, 213 region_name=self.DEFAULT_REGION) 214 215 216def run(): 217 CommandRunner().run() 218 219 220run() 221