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