Add scale command

Closes #9

fig/cli/main.py

@@ -10,6 +10,7 @@ from inspect import getdoc
 
 from .. import __version__
 from ..project import NoSuchService
+from ..service import CannotBeScaledError
 from .command import Command
 from .formatter import Formatter
 from .log_printer import LogPrinter
@@ -82,6 +83,7 @@ class TopLevelCommand(Command):
       ps        List containers
       rm        Remove stopped containers
       run       Run a one-off command
+      scale     Set number of containers for a service
       start     Start services
       stop      Stop services
       up        Create and start containers
@@ -220,6 +222,31 @@ class TopLevelCommand(Command):
                 service.start_container(container, ports=None)
                 c.run()
 
+    def scale(self, options):
+        """
+        Set number of containers to run for a service.
+
+        Numbers are specified in the form `service=num` as arguments.
+        For example:
+
+            $ fig scale web=2 worker=3
+
+        Usage: scale [SERVICE=NUM...]
+        """
+        for s in options['SERVICE=NUM']:
+            if '=' not in s:
+                raise UserError('Arguments to scale should be in the form service=num')
+            service_name, num = s.split('=', 1)
+            try:
+                num = int(num)
+            except ValueError:
+                raise UserError('Number of containers for service "%s" is not a number' % service)
+            try:
+                self.project.get_service(service_name).scale(num)
+            except CannotBeScaledError:
+                raise UserError('Service "%s" cannot be scaled because it specifies a port on the host. If multiple containers for this service were created, the port would clash.\n\nRemove the ":" from the port definition in fig.yml so Docker can choose a random port for each container.' % service_name)
+
+
     def start(self, options):
         """
         Start existing containers.

fig/service.py

@@ -14,6 +14,10 @@ class BuildError(Exception):
     pass
 
 
+class CannotBeScaledError(Exception):
+    pass
+
+
 class Service(object):
     def __init__(self, name, client=None, project='default', links=[], **options):
         if not re.match('^[a-zA-Z0-9]+$', name):
@@ -56,6 +60,40 @@ class Service(object):
             log.info("Killing %s..." % c.name)
             c.kill(**options)
 
+    def scale(self, desired_num):
+        if not self.can_be_scaled():
+            raise CannotBeScaledError()
+
+        # Create enough containers
+        containers = self.containers(stopped=True)
+        while len(containers) < desired_num:
+            containers.append(self.create_container())
+
+        running_containers = []
+        stopped_containers = []
+        for c in containers:
+            if c.is_running:
+                running_containers.append(c)
+            else:
+                stopped_containers.append(c)
+        running_containers.sort(key=lambda c: c.number)
+        stopped_containers.sort(key=lambda c: c.number)
+
+        # Stop containers
+        while len(running_containers) > desired_num:
+            c = running_containers.pop()
+            log.info("Stopping %s..." % c.name)
+            c.stop(timeout=1)
+            stopped_containers.append(c)
+
+        # Start containers
+        while len(running_containers) < desired_num:
+            c = stopped_containers.pop(0)
+            log.info("Starting %s..." % c.name)
+            c.start()
+            running_containers.append(c)
+
+
     def remove_stopped(self, **options):
         for c in self.containers(stopped=True):
             if not c.is_running:
@@ -231,6 +269,12 @@ class Service(object):
         """
         return '%s_%s' % (self.project, self.name)
 
+    def can_be_scaled(self):
+        for port in self.options.get('ports', []):
+            if ':' in str(port):
+                return False
+        return True
+
 
 NAME_RE = re.compile(r'^([^_]+)_([^_]+)_(run_)?(\d+)$')