[benchmarks] add code for stressing the server

2 files changed, 81 insertions, 0 deletions

diff --git a/scripts/scalability/test_controller/server/blobs.py b/scripts/scalability/test_controller/server/blobs.py
new file mode 100644
index 00000000..2a60018d
--- /dev/null
+++ b/scripts/scalability/test_controller/server/blobs.py
@@ -0,0 +1,50 @@
+#!/usr/bin/env python
+
+import shutil
+import os
+
+from argparse import ArgumentParser
+
+from test_controller.utils import mkdir_p, payload
+
+
+def _create_blob(path, data):
+    if not os.path.isfile(path):
+        with open(path, 'w') as f:
+            f.write(data)
+
+
+def create_blobs(target_dir, amount, size):
+    data = payload(size * 1000)
+    for i in xrange(amount):
+        basedir = os.path.join(target_dir, '%d/default/0/0/0' % i)
+        mkdir_p(basedir)
+        _create_blob(os.path.join(basedir, '0'), data)
+
+
+def delete_blobs(target_dir):
+    if not os.path.isdir(target_dir):
+        return
+    for f in os.listdir(target_dir):
+        if f.isdigit():
+            directory = os.path.join(target_dir, f)
+            shutil.rmtree(directory)
+
+
+def parse_args():
+    parser = ArgumentParser()
+    parser.add_argument(
+        'target_dir',
+        help='The target directory where templates will be written to.')
+    parser.add_argument(
+        '--amount', default=1000, type=int,
+        help='The amount of users to create blobs to.')
+    parser.add_argument(
+        '--size', default=1000, type=int,
+        help='The size of each template in KB.')
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    create_blobs(args.target_dir, args.amount, args.size * 1000)
diff --git a/scripts/scalability/test_controller/server/server.tac b/scripts/scalability/test_controller/server/server.tac
index 3d5ea8cc..d5176319 100644
--- a/scripts/scalability/test_controller/server/server.tac
+++ b/scripts/scalability/test_controller/server/server.tac
@@ -43,9 +43,12 @@ import psutil
 from twisted.application import service, internet
 from twisted.web import resource, server
 from twisted.internet.task import LoopingCall
+from twisted.internet.threads import deferToThread
 from twisted.logger import Logger
 
 from test_controller.server.user_dbs import ensure_dbs
+from test_controller.server.blobs import create_blobs
+from test_controller.server.blobs import delete_blobs
 
 
 DEFAULT_HTTP_PORT = 7001
@@ -224,6 +227,33 @@ class SetupResource(resource.Resource):
         request.finish()
 
 
+class BlobsResource(resource.Resource):
+
+    def render_POST(self, request):
+        action = (request.args.get('action') or ['create']).pop()
+        amount = int((request.args.get('amount') or [1000]).pop())
+        size = int((request.args.get('size') or [1000]).pop())
+        if action == 'create':
+            d = deferToThread(create_blobs, '/tmp/soledad-server/blobs',
+                              amount, size)
+        elif action == 'delete':
+            d = deferToThread(delete_blobs, '/tmp/soledad-server/blobs')
+        d.addCallback(self._success, request)
+        d.addErrback(self._error, request)
+        return server.NOT_DONE_YET
+
+    def _success(self, _, request):
+        request.write(SUCCESS)
+        request.finish()
+
+    def _error(self, e, request):
+        message = e.getErrorMessage() if e.getErrorMessage() else repr(e)
+        logger.error('Error processing request: %s' % message)
+        request.setResponseCode(500)
+        request.write(json.dumps({'error': str(e)}))
+        request.finish()
+
+
 class Root(resource.Resource):
 
     def __init__(self):
@@ -231,6 +261,7 @@ class Root(resource.Resource):
         self.putChild('mem', MonitorResource(MemoryWatcher))
         self.putChild('cpu', MonitorResource(CpuWatcher))
         self.putChild('setup', SetupResource())
+        self.putChild('blobs', BlobsResource())
 
 
 application = service.Application("Resource Monitor")