add storage policy option

provide ability for user to select which storage policy
to run benchmark against

Change-Id: I9973f964ec3989281f47f853be053498e7a6e51e
Signed-off-by: Thiago da Silva <thiago@redhat.com>

3 files changed, 18 insertions, 3 deletions

diff --git a/bin/swift-bench b/bin/swift-bench
index a1d7a19..b516c24 100755
--- a/bin/swift-bench
+++ b/bin/swift-bench
@@ -74,6 +74,7 @@ if __name__ == '__main__':
     num_gets = 10000
     delete = yes
     auth_version = 1.0
+    policy_name = gold
     """
     parser = OptionParser(usage=usage)
     parser.add_option('', '--saio', dest='saio', action='store_true',
@@ -120,6 +121,9 @@ if __name__ == '__main__':
                       help='Authentication version')
     parser.add_option('-d', '--delay', dest='delay',
                       help='Delay before delete requests in seconds')
+    parser.add_option('-P', '--policy-name', dest='policy_name',
+                      help='Specify which policy to use when creating '
+                           'containers')
 
     if len(sys.argv) == 1:
         parser.print_help()
diff --git a/etc/swift-bench.conf-sample b/etc/swift-bench.conf-sample
index 423608e..8ca7f8b 100644
--- a/etc/swift-bench.conf-sample
+++ b/etc/swift-bench.conf-sample
@@ -34,6 +34,9 @@
 # The base name for created containers.
 # container_name = (randomly-chosen uuid4)
 
+# The storage policy to use when creating containers
+# policy_name = gold
+
 # Should swift-bench benchmark DELETEing the created objects and then delete
 # all created containers?
 # delete = yes
diff --git a/swiftbench/bench.py b/swiftbench/bench.py
index bc899ea..65f3f87 100644
--- a/swiftbench/bench.py
+++ b/swiftbench/bench.py
@@ -45,13 +45,13 @@ except ImportError:
 HTTP_CONFLICT = 409
 
 
-def _func_on_containers(logger, conf, concurrency_key, func):
+def _func_on_containers(logger, conf, concurrency_key, func, **kwargs):
     """Run a function on each container with concurrency."""
 
     bench = Bench(logger, conf, [])
     pool = eventlet.GreenPool(int(getattr(conf, concurrency_key)))
     for container in conf.containers:
-        pool.spawn_n(func, bench.url, bench.token, container)
+        pool.spawn_n(func, bench.url, bench.token, container, **kwargs)
     pool.waitall()
 
 
@@ -73,7 +73,15 @@ def delete_containers(logger, conf):
 def create_containers(logger, conf):
     """Utility function to create benchmark containers."""
 
-    _func_on_containers(logger, conf, 'put_concurrency', client.put_container)
+    if conf.policy_name:
+        logger.info("Creating containers with storage policy: %s" %
+                    conf.policy_name)
+        _func_on_containers(logger, conf, 'put_concurrency',
+                            client.put_container,
+                            headers={'X-Storage-Policy': conf.policy_name})
+    else:
+        _func_on_containers(logger, conf, 'put_concurrency',
+                            client.put_container)
 
 
 class SourceFile(object):