Merge lp:~fwereade/pyjuju/agent-trash-old-sessions into lp:pyjuju

[William Reade]

Agents now require a --session-file argument, which is written to on connect and deleted on close.

If, on startup, the session file exists and refers to a valid session, the session is explicitly closed before any further work is done.

IMO this is the best approach because it minimises the number of cases we have to think about; rather than attempting to re-engage with an existing session (set up watches but skip ephemeral node creation... more considerations?), we always have to set up exactly the same state, whatever the reason we're restarting.

Session files are chmodded to 600 but, on reflection, I'm not sure this actually helps... a malicious hook could still plausibly hijack them. Thoughts?

The cloud-init runcmds have changed again, and these make up a tediously large proportion of the diff. Sorry for the noise :(.

[Kapil Thangavelu (hazmat)]

Excerpts from William Reade's message of Wed Nov 30 09:56:37 UTC 2011:
> William Reade has proposed merging lp:~fwereade/juju/agent-trash-old-sessions into lp:juju with lp:~fwereade/juju/upstartify-agents as a prerequisite.
>
> Requested reviews:
> juju hackers (juju)
> Related bugs:
> Bug #898082 in juju: "immediate agent restart fails"
> https://bugs.launchpad.net/juju/+bug/898082
>
> For more details, see:
> https://code.launchpad.net/~fwereade/juju/agent-trash-old-sessions/+merge/83905
>
> Agents now require a --session-file argument, which is written to on connect and deleted on close.
>
> If, on startup, the session file exists and refers to a valid session, the session is explicitly closed before any further work is done.
>
> IMO this is the best approach because it minimises the number of cases we have to think about; rather than attempting to re-engage with an existing session (set up watches but skip ephemeral node creation... more considerations?), we always have to set up exactly the same state, whatever the reason we're restarting.
>
> Session files are chmodded to 600 but, on reflection, I'm not sure this actually helps... a malicious hook could still plausibly hijack them. Thoughts?
>
> The cloud-init runcmds have changed again, and these make up a tediously large proportion of the diff. Sorry for the noise :(.

Haven't looked at the diff, just a comment on the approach, this seems pretty
reasonable and simple. The watch callbacks (transient registration with libzk)
are gone so we can't usefully reconnect to the original session watches, we'd
have to restablish them anyways. Effectively sessions are primarily useful for
transient error activity not process restarts.

However there is a caveat wrt to session closing and ephemeral nodes that's
worth keeping in mind, namely that the ephemerals are coordination points for
things like relation membership notification, and this will result in transient
-depart/-join hook execution as the old session closes and the new one comes
into place. I don't see any reasonable alternatives that are significantly more
complex though.

[lp:~fwereade/pyjuju/agent-trash-old-sessions]

476. By William Reade on 2011-12-02

merge parent

[William Reade (fwereade)]

rejecting due to abandonment of prereq

[lp:~fwereade/pyjuju/agent-trash-old-sessions]

477. By William Reade on 2011-12-12

agents now require --session-file, while will be used to shoot old sessions in the head when the agent comes up, to ensure we're starting from clean ZK state

478. By William Reade on 2011-12-12

merge parent

479. By William Reade on 2011-12-15

merge parent

480. By William Reade on 2011-12-16

merge parent

481. By William Reade on 2012-01-10

merge parent

482. By William Reade on 2012-01-11

merge parent

483. By William Reade on 2012-01-12

merge parent

484. By William Reade on 2012-01-13

merge parent

485. By William Reade on 2012-01-13

address review point

486. By William Reade on 2012-01-13

merge parent

487. By William Reade on 2012-01-16

merge parent

488. By William Reade on 2012-01-16

merge parent

489. By William Reade on 2012-01-23

merge parent

490. By William Reade on 2012-01-27

merge parent

491. By William Reade on 2012-02-02

merge parent

492. By William Reade on 2012-02-08

merge parent

493. By William Reade on 2012-02-21

merge parent

Unmerged revisions

=== modified file 'examples/oneiric/mysql/hooks/install'
--- examples/oneiric/mysql/hooks/install	2011-09-15 18:56:08 +0000
+++ examples/oneiric/mysql/hooks/install	2011-12-02 14:44:25 +0000
@@ -21,5 +21,5 @@
 juju-log "Editing my.cnf to allow listening on all interfaces"
 sed --in-place=old 's/127\.0\.0\.1/0.0.0.0/' /etc/mysql/my.cnf
 
-juju-log "Restarting mysql service"
-service mysql restart
+juju-log "Stopping mysql service"
+service mysql stop

=== modified file 'examples/oneiric/mysql/hooks/start'
--- examples/oneiric/mysql/hooks/start	2011-02-03 01:23:43 +0000
+++ examples/oneiric/mysql/hooks/start	2011-12-02 14:44:25 +0000
@@ -1,1 +1,3 @@
-#!/bin/bash
\ No newline at end of file
+#!/bin/bash
+juju-log "Starting mysql service"
+service mysql start || service mysql restart

=== modified file 'examples/oneiric/mysql/hooks/stop'
--- examples/oneiric/mysql/hooks/stop	2011-09-15 18:56:08 +0000
+++ examples/oneiric/mysql/hooks/stop	2011-12-02 14:44:25 +0000
@@ -1,3 +1,3 @@
 #!/bin/bash
 juju-log "Stopping mysql service"
-/etc/init.d/mysql stop
+service mysql stop || true

=== modified file 'juju/agents/base.py'
--- juju/agents/base.py	2011-09-22 13:23:00 +0000
+++ juju/agents/base.py	2011-12-02 14:44:25 +0000
@@ -1,7 +1,9 @@
 import argparse
 import os
+import logging
+import stat
 import sys
-import logging
+import yaml
 
 import zookeeper
 
@@ -18,6 +20,23 @@
 from juju.state.environment import GlobalSettingsStateManager
 
 
+def load_client_id(path):
+    try:
+        with open(path) as f:
+            return yaml.load(f.read())
+    except IOError:
+        return None
+
+
+def save_client_id(path, client_id):
+    parent = os.path.dirname(path)
+    if not os.path.exists(parent):
+        os.makedirs(parent)
+    with open(path, "w") as f:
+        f.write(yaml.dump(client_id))
+    os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)
+
+
 class TwistedOptionNamespace(object):
     """
     An argparse namespace implementation that is compatible with twisted
@@ -153,13 +172,40 @@
                 "Invalid juju-directory %r, does not exist." % (
                     options.get("juju_directory")))
 
+        if options["session_file"] is None:
+            raise JujuError("No session file specified")
+
         self.config = options
 
     @inlineCallbacks
+    def _kill_existing_session(self):
+        try:
+            # We might have died suddenly, in which case the session may
+            # still be alive. If this is the case, shoot it in the head, so
+            # it doesn't interfere with our attempts to recreate our state.
+            # (We need to be able to recreate our state *anyway*, and it's
+            # much simpler to force ourselves to recreate it every time than
+            # it is to mess around partially recreating partial state.)
+            client_id = load_client_id(self.config["session_file"])
+            if client_id is None:
+                return
+            temp_client = yield ZookeeperClient().connect(
+                self.config["zookeeper_servers"], client_id=client_id)
+            yield temp_client.close()
+        except zookeeper.ZooKeeperException:
+            # We don't really care what went wrong; just that we're not able
+            # to connect using the old session, and therefore we should be ok
+            # to start a fresh one without transient state hanging around.
+            pass
+
+    @inlineCallbacks
     def connect(self):
         """Return an authenticated connection to the juju zookeeper."""
-        hosts = self.config["zookeeper_servers"]
-        self.client = yield ZookeeperClient().connect(hosts)
+        yield self._kill_existing_session()
+        self.client = yield ZookeeperClient().connect(
+            self.config["zookeeper_servers"])
+        save_client_id(
+            self.config["session_file"], self.client.client_id)
 
         principals = self.config.get("principals", ())
         for principal in principals:
@@ -200,6 +246,9 @@
         finally:
             if self.client and self.client.connected:
                 self.client.close()
+            session_file = self.config["session_file"]
+            if os.path.exists(session_file):
+                os.unlink(session_file)
 
     def set_watch_enabled(self, flag):
         """Set boolean flag for whether this agent should watching zookeeper.
@@ -285,3 +334,7 @@
     parser.add_argument(
         "--juju-directory", default=juju_home, type=os.path.abspath,
         help="juju working directory ($JUJU_HOME)")
+
+    parser.add_argument(
+        "--session-file", default=None, type=os.path.abspath,
+        help="like a pidfile, but for the zookeeper session id")

=== modified file 'juju/agents/tests/common.py'
--- juju/agents/tests/common.py	2011-09-15 19:24:47 +0000
+++ juju/agents/tests/common.py	2011-12-02 14:44:25 +0000
@@ -55,6 +55,7 @@
         options = TwistedOptionNamespace()
         options["juju_directory"] = self.juju_directory
         options["zookeeper_servers"] = get_test_zookeeper_address()
+        options["session_file"] = self.makeFile()
         return succeed(options)
 
     @inlineCallbacks

=== modified file 'juju/agents/tests/test_base.py'
--- juju/agents/tests/test_base.py	2011-12-02 14:44:24 +0000
+++ juju/agents/tests/test_base.py	2011-12-02 14:44:25 +0000
@@ -2,13 +2,14 @@
 import json
 import logging
 import os
+import stat
 import sys
-
+import yaml
 
 from twisted.application.app import AppLogger
 from twisted.application.service import IService, IServiceCollection
 from twisted.internet.defer import (
-    succeed, Deferred, inlineCallbacks, returnValue)
+    fail, succeed, Deferred, inlineCallbacks, returnValue)
 from twisted.python.components import Componentized
 from twisted.python import log
 
@@ -20,6 +21,7 @@
 
 from juju.agents.base import (
     BaseAgent, TwistedOptionNamespace, AgentRunner, AgentLogger)
+from juju.agents.dummy import DummyAgent
 from juju.errors import NoConnection, JujuError
 from juju.lib.zklog import ZookeeperHandler
 
@@ -34,8 +36,8 @@
     @inlineCallbacks
     def setUp(self):
         yield super(BaseAgentTest, self).setUp()
-        self.change_environment(
-            JUJU_HOME=self.makeDir())
+        self.juju_home = self.makeDir()
+        self.change_environment(JUJU_HOME=self.juju_home)
 
     def test_as_app(self):
         """The agent class can be accessed as an application."""
@@ -79,6 +81,8 @@
         # Agent options
         self.assertEqual(parser.get_default("principals"), [])
         self.assertEqual(parser.get_default("zookeeper_servers"), "")
+        self.assertEqual(parser.get_default("juju_directory"), self.juju_home)
+        self.assertEqual(parser.get_default("session_file"), None)
 
     def test_twistd_flags_correspond(self):
         parser = argparse.ArgumentParser()
@@ -99,7 +103,8 @@
         log_file_path = self.makeFile()
 
         options = parser.parse_args(
-            ["--logfile", log_file_path], namespace=TwistedOptionNamespace())
+            ["--logfile", log_file_path, "--session-file", self.makeFile()],
+            namespace=TwistedOptionNamespace())
 
         def match_observer(observer):
             return isinstance(observer.im_self, log.PythonLoggingObserver)
@@ -180,8 +185,9 @@
         This will create an agent instance,  parse the cli args, passes them to
         the agent, and starts the agent runner.
         """
-        self.change_args("es-agent", "--zookeeper-servers",
-                         get_test_zookeeper_address())
+        self.change_args(
+            "es-agent", "--zookeeper-servers", get_test_zookeeper_address(),
+            "--session-file", self.makeFile())
         runner = self.mocker.patch(AgentRunner)
         runner.run()
         mock_agent = self.mocker.patch(BaseAgent)
@@ -220,7 +226,8 @@
 
         self.change_args(
             "es-agent", "--nodaemon",
-            "--zookeeper-servers", get_test_zookeeper_address())
+            "--zookeeper-servers", get_test_zookeeper_address(),
+            "--session-file", self.makeFile())
         runner = self.mocker.patch(AgentRunner)
         logger = self.mocker.patch(AppLogger)
         logger.start(MATCH_APP)
@@ -230,6 +237,7 @@
         DummyAgent.run()
         return started
 
+    @inlineCallbacks
     def test_stop_service_stub_closes_agent(self):
         """The base class agent, stopService will the stop method.
 
@@ -238,6 +246,7 @@
         """
         mock_agent = self.mocker.patch(BaseAgent)
         mock_client = self.mocker.mock(ZookeeperClient)
+        session_file = self.makeFile()
 
         # connection is closed after agent.stop invoked.
         with self.mocker.order():
@@ -259,11 +268,17 @@
             self.mocker.result(mock_client)
             mock_client.close()
 
+            # delete session file
+            mock_agent.config
+            self.mocker.result({"session_file": session_file})
+
         self.mocker.replay()
 
         agent = BaseAgent()
-        return agent.stopService()
+        yield agent.stopService()
+        self.assertFalse(os.path.exists(session_file))
 
+    @inlineCallbacks
     def test_stop_service_stub_ignores_disconnected_agent(self):
         """The base class agent, stopService will the stop method.
 
@@ -271,6 +286,7 @@
         """
         mock_agent = self.mocker.patch(BaseAgent)
         mock_client = self.mocker.mock(ZookeeperClient)
+        session_file = self.makeFile()
 
         # connection is closed after agent.stop invoked.
         with self.mocker.order():
@@ -286,10 +302,14 @@
             mock_client.connected
             self.mocker.result(False)
 
+            mock_agent.config
+            self.mocker.result({"session_file": session_file})
+
         self.mocker.replay()
 
         agent = BaseAgent()
-        return agent.stopService()
+        yield agent.stopService()
+        self.assertFalse(os.path.exists(session_file))
 
     def test_run_base_raises_error(self):
         """The base class agent, raises a notimplemented error when started."""
@@ -297,12 +317,15 @@
         client.connect(get_test_zookeeper_address())
         client_mock = self.mocker.mock()
         self.mocker.result(succeed(client_mock))
+        client_mock.client_id
+        self.mocker.result((123, "abc"))
         self.mocker.replay()
 
         agent = BaseAgent()
         agent.configure({
             "zookeeper_servers": get_test_zookeeper_address(),
-            "juju_directory": self.makeDir()})
+            "juju_directory": self.makeDir(),
+            "session_file": self.makeFile()})
         d = agent.startService()
         self.failUnlessFailure(d, NotImplementedError)
         return d
@@ -313,22 +336,39 @@
         client = self.mocker.patch(ZookeeperClient)
         client.connect("x2.example.com")
         self.mocker.result(succeed(mock_client))
+        mock_client.client_id
+        self.mocker.result((123, "abc"))
         self.mocker.replay()
 
         agent = BaseAgent()
         agent.configure({"zookeeper_servers": "x2.example.com",
-                         "juju_directory": self.makeDir()})
+                         "juju_directory": self.makeDir(),
+                         "session_file": self.makeFile()})
         result = agent.connect()
         self.assertEqual(result.result, mock_client)
         self.assertEqual(agent.client, mock_client)
 
-    def test_non_existant_directory(self):
+    def test_nonexistent_directory(self):
         """If the juju directory does not exist an error should be raised.
         """
         juju_directory = self.makeDir()
         os.rmdir(juju_directory)
         data = {"zookeeper_servers": get_test_zookeeper_address(),
-                "juju_directory": juju_directory}
+                "juju_directory": juju_directory,
+                "session_file": self.makeFile()}
+
+        agent = BaseAgent()
+        self.assertRaises(
+            JujuError,
+            agent.configure,
+            data)
+
+    def test_bad_session_file(self):
+        """If the session file cannot be created an error should be raised.
+        """
+        data = {"zookeeper_servers": get_test_zookeeper_address(),
+                "juju_directory": self.makeDir(),
+                "session_file": None}
 
         agent = BaseAgent()
         self.assertRaises(
@@ -341,7 +381,8 @@
         juju_directory = self.makeDir()
         self.change_args(
             "es-agent", "--zookeeper-servers", get_test_zookeeper_address(),
-            "--juju-directory", juju_directory)
+            "--juju-directory", juju_directory,
+            "--session-file", self.makeFile())
 
         agent = BaseAgent()
         parser = argparse.ArgumentParser()
@@ -363,7 +404,9 @@
         agent = BaseAgent()
         parser = argparse.ArgumentParser()
         agent.setup_options(parser)
-        options = parser.parse_args(namespace=TwistedOptionNamespace())
+        options = parser.parse_args(
+            ["--session-file", self.makeFile()],
+            namespace=TwistedOptionNamespace())
         agent.configure(options)
         self.assertEqual(
             agent.config["juju_directory"], juju_directory)
@@ -379,6 +422,8 @@
         client = self.mocker.patch(ZookeeperClient)
         client.connect("x1.example.com")
         self.mocker.result(succeed(client))
+        client.client_id
+        self.mocker.result((123, "abc"))
         client.add_auth("digest", "admin:abc")
         client.add_auth("digest", "agent:xyz")
         client.exists("/")
@@ -387,7 +432,105 @@
         agent = BaseAgent()
         parser = argparse.ArgumentParser()
         agent.setup_options(parser)
-        options = parser.parse_args(namespace=TwistedOptionNamespace())
+        options = parser.parse_args(
+            ["--session-file", self.makeFile()],
+            namespace=TwistedOptionNamespace())
+        agent.configure(options)
+        d = agent.startService()
+        self.failUnlessFailure(d, NotImplementedError)
+        return d
+
+    def test_connect_closes_running_session(self):
+        self.change_args("es-agent")
+        self.change_environment(
+            JUJU_HOME=self.makeDir(),
+            JUJU_ZOOKEEPER="x1.example.com")
+
+        session_file = self.makeFile()
+        with open(session_file, "w") as f:
+            f.write(yaml.dump((123, "abc")))
+        mock_client_1 = self.mocker.mock()
+        client = self.mocker.patch(ZookeeperClient)
+        client.connect("x1.example.com", client_id=(123, "abc"))
+        self.mocker.result(succeed(mock_client_1))
+        mock_client_1.close()
+        self.mocker.result(None)
+
+        mock_client_2 = self.mocker.mock()
+        client.connect("x1.example.com")
+        self.mocker.result(succeed(mock_client_2))
+        mock_client_2.client_id
+        self.mocker.result((456, "def"))
+        self.mocker.replay()
+
+        agent = BaseAgent()
+        parser = argparse.ArgumentParser()
+        agent.setup_options(parser)
+        options = parser.parse_args(
+            ["--session-file", session_file],
+            namespace=TwistedOptionNamespace())
+        agent.configure(options)
+        d = agent.startService()
+        self.failUnlessFailure(d, NotImplementedError)
+        return d
+
+    def test_connect_handles_expired_session(self):
+        self.change_args("es-agent")
+        self.change_environment(
+            JUJU_HOME=self.makeDir(),
+            JUJU_ZOOKEEPER="x1.example.com")
+
+        session_file = self.makeFile()
+        with open(session_file, "w") as f:
+            f.write(yaml.dump((123, "abc")))
+        client = self.mocker.patch(ZookeeperClient)
+        client.connect("x1.example.com", client_id=(123, "abc"))
+        self.mocker.result(fail(zookeeper.SessionExpiredException()))
+
+        mock_client = self.mocker.mock()
+        client.connect("x1.example.com")
+        self.mocker.result(succeed(mock_client))
+        mock_client.client_id
+        self.mocker.result((456, "def"))
+        self.mocker.replay()
+
+        agent = BaseAgent()
+        parser = argparse.ArgumentParser()
+        agent.setup_options(parser)
+        options = parser.parse_args(
+            ["--session-file", session_file],
+            namespace=TwistedOptionNamespace())
+        agent.configure(options)
+        d = agent.startService()
+        self.failUnlessFailure(d, NotImplementedError)
+        return d
+
+    def test_connect_handles_nonsense_session(self):
+        self.change_args("es-agent")
+        self.change_environment(
+            JUJU_HOME=self.makeDir(),
+            JUJU_ZOOKEEPER="x1.example.com")
+
+        session_file = self.makeFile()
+        with open(session_file, "w") as f:
+            f.write(yaml.dump("cheesy wotsits"))
+        client = self.mocker.patch(ZookeeperClient)
+        client.connect("x1.example.com", client_id="cheesy wotsits")
+        self.mocker.result(fail(zookeeper.ZooKeeperException()))
+
+        mock_client = self.mocker.mock()
+        client.connect("x1.example.com")
+        self.mocker.result(succeed(mock_client))
+        mock_client.client_id
+        self.mocker.result((456, "def"))
+        self.mocker.replay()
+
+        agent = BaseAgent()
+        parser = argparse.ArgumentParser()
+        agent.setup_options(parser)
+        options = parser.parse_args(
+            ["--session-file", session_file],
+            namespace=TwistedOptionNamespace())
         agent.configure(options)
         d = agent.startService()
         self.failUnlessFailure(d, NotImplementedError)
@@ -405,6 +548,21 @@
         agent.set_watch_enabled(False)
         self.assertFalse(agent.get_watch_enabled())
 
+    @inlineCallbacks
+    def test_session_file_permissions(self):
+        session_file = self.makeFile()
+        agent = DummyAgent()
+        agent.configure({
+            "session_file": session_file,
+            "juju_directory": self.makeDir(),
+            "zookeeper_servers": get_test_zookeeper_address()})
+        yield agent.startService()
+        mode = os.stat(session_file).st_mode
+        mask = stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO
+        self.assertEquals(mode & mask, stat.S_IRUSR | stat.S_IWUSR)
+        yield agent.stopService()
+        self.assertFalse(os.path.exists(session_file))
+
 
 class AgentDebugLogSettingsWatch(AgentTestBase):
 

=== modified file 'juju/agents/tests/test_machine.py'
--- juju/agents/tests/test_machine.py	2011-10-05 13:59:44 +0000
+++ juju/agents/tests/test_machine.py	2011-12-02 14:44:25 +0000
@@ -120,10 +120,9 @@
         # initially setup by get_agent_config in setUp
         self.change_environment(JUJU_MACHINE_ID="")
         self.change_args("es-agent",
-                         "--zookeeper-servers",
-                         get_test_zookeeper_address(),
-                         "--juju-directory",
-                         self.makeDir())
+                         "--zookeeper-servers", get_test_zookeeper_address(),
+                         "--juju-directory", self.makeDir(),
+                         "--session-file", self.makeFile())
         parser = argparse.ArgumentParser()
         self.agent.setup_options(parser)
         options = parser.parse_args(namespace=TwistedOptionNamespace())

=== modified file 'juju/agents/tests/test_unit.py'
--- juju/agents/tests/test_unit.py	2011-12-02 14:44:24 +0000
+++ juju/agents/tests/test_unit.py	2011-12-02 14:44:25 +0000
@@ -170,7 +170,8 @@
         self.change_args(
             "unit-agent",
             "--juju-directory", self.makeDir(),
-            "--zookeeper-servers", get_test_zookeeper_address())
+            "--zookeeper-servers", get_test_zookeeper_address(),
+            "--session-file", self.makeFile())
 
         parser = argparse.ArgumentParser()
         self.agent.setup_options(parser)
@@ -202,6 +203,7 @@
         options = {}
         options["juju_directory"] = self.juju_directory
         options["zookeeper_servers"] = get_test_zookeeper_address()
+        options["session_file"] = self.makeFile()
         options["unit_name"] = "rabbit-1"
         agent = self.agent_class()
         agent.configure(options)

=== modified file 'juju/agents/unit.py'
--- juju/agents/unit.py	2011-12-02 14:44:24 +0000
+++ juju/agents/unit.py	2011-12-02 14:44:25 +0000
@@ -80,10 +80,11 @@
             self.config["juju_directory"], "state")
 
         # Setup the server portion of the cli api exposed to hooks.
+        socket_path = os.path.join(self.unit_directory, HOOK_SOCKET_FILE)
+        if os.path.exists(socket_path):
+            os.unlink(socket_path)
         from twisted.internet import reactor
-        self.api_socket = reactor.listenUNIX(
-            os.path.join(self.unit_directory, HOOK_SOCKET_FILE),
-            self.api_factory)
+        self.api_socket = reactor.listenUNIX(socket_path, self.api_factory)
 
         # Setup the unit state's address
         address = yield get_unit_address(self.client)
@@ -98,6 +99,10 @@
         # Inform the system, we're alive.
         yield self.unit_state.connect_agent()
 
+        # Start paying attention to the debug-log setting
+        if self.get_watch_enabled():
+            yield self.unit_state.watch_hook_debug(self.cb_watch_hook_debug)
+
         self.lifecycle = UnitLifecycle(
             self.client,
             self.unit_state,

=== modified file 'juju/control/tests/test_status.py'
--- juju/control/tests/test_status.py	2011-11-30 00:57:05 +0000
+++ juju/control/tests/test_status.py	2011-12-02 14:44:25 +0000
@@ -107,6 +107,7 @@
         options = TwistedOptionNamespace()
         options["juju_directory"] = path
         options["zookeeper_servers"] = get_test_zookeeper_address()
+        options["session_file"] = self.makeFile()
         for k, v in extra_options.items():
             options[k] = v
         agent.configure(options)
@@ -302,6 +303,7 @@
         options = TwistedOptionNamespace()
         options["juju_directory"] = self.makeDir()
         options["zookeeper_servers"] = get_test_zookeeper_address()
+        options["session_file"] = self.makeFile()
         options["machine_id"] = "0"
         agent.configure(options)
         agent.set_watch_enabled(False)

=== modified file 'juju/lib/statemachine.py'
--- juju/lib/statemachine.py	2011-05-04 19:40:59 +0000
+++ juju/lib/statemachine.py	2011-12-02 14:44:25 +0000
@@ -213,7 +213,7 @@
         state_dict = yield self._load()
         if not state_dict:
             returnValue({})
-        returnValue(state_dict["state_variables"])
+        returnValue(state_dict.get("state_variables", {}))
 
     def set_observer(self, observer):
         """Set a callback, that will be notified on state changes.

=== modified file 'juju/machine/tests/test_unit_deployment.py'
--- juju/machine/tests/test_unit_deployment.py	2011-12-02 14:44:24 +0000
+++ juju/machine/tests/test_unit_deployment.py	2011-12-02 14:44:25 +0000
@@ -114,7 +114,9 @@
                 self.deployment.directory, "charm.log")
             command = " ".join([
                 sys.executable, "-m", "juju.agents.dummy",
-                "--nodaemon", "--logfile", log_file])
+                "--nodaemon", "--logfile", log_file,
+                "--session-file",
+                "/var/run/juju/unit-wordpress-0-agent.zksession"])
             self.assertEquals(exec_, command)
         d.addCallback(verify_upstart)
         return d

=== modified file 'juju/machine/unit.py'
--- juju/machine/unit.py	2011-12-02 14:44:24 +0000
+++ juju/machine/unit.py	2011-12-02 14:44:25 +0000
@@ -87,7 +87,9 @@
         self.service.set_command(" ".join((
             sys.executable, "-m", self.unit_agent_module,
             "--nodaemon",
-            "--logfile", os.path.join(self.directory, "charm.log"))))
+            "--logfile", os.path.join(self.directory, "charm.log"),
+            "--session-file",
+            "/var/run/juju/unit-%s-agent.zksession" % self.unit_path_name)))
         return self.service.start()
 
     @inlineCallbacks
@@ -230,7 +232,9 @@
             sys.executable,
             "-m", "juju.agents.unit",
             "--nodaemon",
-            "--logfile", "/var/log/juju/unit-%s.log" % self.unit_path_name)))
+            "--logfile", "/var/log/juju/unit-%s.log" % self.unit_path_name,
+            "--session-file",
+            "/var/run/juju/unit-%s-agent.zksession" % self.unit_path_name)))
         yield service.install()
 
         # Create a symlink on the host for easier access to the unit log file

=== modified file 'juju/providers/common/cloudinit.py'
--- juju/providers/common/cloudinit.py	2011-12-02 14:44:24 +0000
+++ juju/providers/common/cloudinit.py	2011-12-02 14:44:25 +0000
@@ -48,7 +48,8 @@
         {"JUJU_MACHINE_ID": machine_id, "JUJU_ZOOKEEPER": zookeeper_hosts})
     service.set_command(
         "python -m juju.agents.machine --nodaemon "
-        "--logfile /var/log/juju/machine-agent.log")
+        "--logfile /var/log/juju/machine-agent.log "
+        "--session-file /var/run/juju/machine-agent.zksession")
     return service.get_cloud_init_commands()
 
 
@@ -58,7 +59,8 @@
     service.set_environ({"JUJU_ZOOKEEPER": zookeeper_hosts})
     service.set_command(
         "python -m juju.agents.provision --nodaemon "
-        "--logfile /var/log/juju/provision-agent.log")
+        "--logfile /var/log/juju/provision-agent.log "
+        "--session-file /var/run/juju/provision-agent.zksession")
     return service.get_cloud_init_commands()
 
 

=== modified file 'juju/providers/common/tests/data/cloud_init_bootstrap'
--- juju/providers/common/tests/data/cloud_init_bootstrap	2011-12-02 14:44:24 +0000
+++ juju/providers/common/tests/data/cloud_init_bootstrap	2011-12-02 14:44:25 +0000
@@ -28,6 +28,7 @@
 
 
     exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession
 
     EOF
 
@@ -50,6 +51,7 @@
 
 
     exec python -m juju.agents.provision --nodaemon --logfile /var/log/juju/provision-agent.log
+    --session-file /var/run/juju/provision-agent.zksession
 
     EOF
 

=== modified file 'juju/providers/common/tests/data/cloud_init_bootstrap_zookeepers'
--- juju/providers/common/tests/data/cloud_init_bootstrap_zookeepers	2011-12-02 14:44:24 +0000
+++ juju/providers/common/tests/data/cloud_init_bootstrap_zookeepers	2011-12-02 14:44:25 +0000
@@ -28,6 +28,7 @@
 
 
     exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession
 
     EOF
 
@@ -50,6 +51,7 @@
 
 
     exec python -m juju.agents.provision --nodaemon --logfile /var/log/juju/provision-agent.log
+    --session-file /var/run/juju/provision-agent.zksession
 
     EOF
 

=== modified file 'juju/providers/common/tests/data/cloud_init_branch'
--- juju/providers/common/tests/data/cloud_init_branch	2011-12-02 14:44:24 +0000
+++ juju/providers/common/tests/data/cloud_init_branch	2011-12-02 14:44:25 +0000
@@ -31,6 +31,7 @@
 
 
     exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession
 
     EOF
 

=== modified file 'juju/providers/common/tests/data/cloud_init_branch_trunk'
--- juju/providers/common/tests/data/cloud_init_branch_trunk	2011-12-02 14:44:24 +0000
+++ juju/providers/common/tests/data/cloud_init_branch_trunk	2011-12-02 14:44:25 +0000
@@ -31,6 +31,7 @@
 
 
     exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession
 
     EOF
 

=== modified file 'juju/providers/common/tests/data/cloud_init_distro'
--- juju/providers/common/tests/data/cloud_init_distro	2011-12-02 14:44:24 +0000
+++ juju/providers/common/tests/data/cloud_init_distro	2011-12-02 14:44:25 +0000
@@ -27,6 +27,7 @@
 
 
     exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession
 
     EOF
 

=== modified file 'juju/providers/common/tests/data/cloud_init_ppa'
--- juju/providers/common/tests/data/cloud_init_ppa	2011-12-02 14:44:24 +0000
+++ juju/providers/common/tests/data/cloud_init_ppa	2011-12-02 14:44:25 +0000
@@ -29,6 +29,7 @@
 
 
     exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession
 
     EOF
 

=== modified file 'juju/providers/ec2/tests/data/bootstrap_cloud_init'
--- juju/providers/ec2/tests/data/bootstrap_cloud_init	2011-12-02 14:44:24 +0000
+++ juju/providers/ec2/tests/data/bootstrap_cloud_init	2011-12-02 14:44:25 +0000
@@ -28,6 +28,7 @@
 
 
     exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession
 
     EOF
 
@@ -50,6 +51,7 @@
 
 
     exec python -m juju.agents.provision --nodaemon --logfile /var/log/juju/provision-agent.log
+    --session-file /var/run/juju/provision-agent.zksession
 
     EOF
 

=== modified file 'juju/providers/ec2/tests/data/launch_cloud_init'
--- juju/providers/ec2/tests/data/launch_cloud_init	2011-12-02 14:44:24 +0000
+++ juju/providers/ec2/tests/data/launch_cloud_init	2011-12-02 14:44:25 +0000
@@ -27,6 +27,7 @@
 
 
     exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession
 
     EOF
 

=== modified file 'juju/providers/ec2/tests/data/launch_cloud_init_branch'
--- juju/providers/ec2/tests/data/launch_cloud_init_branch	2011-12-02 14:44:24 +0000
+++ juju/providers/ec2/tests/data/launch_cloud_init_branch	2011-12-02 14:44:25 +0000
@@ -31,6 +31,7 @@
 
 
     exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession
 
     EOF
 

=== modified file 'juju/providers/ec2/tests/data/launch_cloud_init_ppa'
--- juju/providers/ec2/tests/data/launch_cloud_init_ppa	2011-12-02 14:44:24 +0000
+++ juju/providers/ec2/tests/data/launch_cloud_init_ppa	2011-12-02 14:44:25 +0000
@@ -29,6 +29,7 @@
 
 
     exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession
 
     EOF
 

=== modified file 'juju/providers/local/agent.py'
--- juju/providers/local/agent.py	2011-12-02 14:44:24 +0000
+++ juju/providers/local/agent.py	2011-12-02 14:44:25 +0000
@@ -47,7 +47,8 @@
         self._service.set_environ(env)
         self._service_args = [
             sys.executable, "-m", self.agent_module,
-            "--nodaemon", "--logfile", log_file]
+            "--nodaemon", "--logfile", log_file,
+            "--session-file", "/var/run/juju/%s-machine-agent.zksession" % juju_unit_namespace]
 
     @property
     def juju_origin(self):

=== modified file 'juju/providers/local/tests/test_agent.py'
--- juju/providers/local/tests/test_agent.py	2011-12-02 14:44:24 +0000
+++ juju/providers/local/tests/test_agent.py	2011-12-02 14:44:25 +0000
@@ -14,7 +14,6 @@
 
 class ManagedAgentTest(TestCase):
 
-
     @inlineCallbacks
     def test_managed_agent_config(self):
         subprocess_calls = []
@@ -60,8 +59,10 @@
             if line.startswith("exec"):
                 exec_ = line[5:-1]
 
-        expect_exec = "%s -m juju.agents.machine --nodaemon --logfile %s" % (
-            sys.executable, log_file)
+        expect_exec = (
+            "%s -m juju.agents.machine --nodaemon --logfile %s "
+            "--session-file /var/run/juju/ns1-machine-agent.zksession"
+            % (sys.executable, log_file))
         self.assertEquals(exec_, expect_exec)
 
         env = dict((k, v.strip('"')) for (k, v) in env)

=== modified file 'juju/providers/orchestra/tests/data/bootstrap_user_data'
--- juju/providers/orchestra/tests/data/bootstrap_user_data	2011-12-02 14:44:24 +0000
+++ juju/providers/orchestra/tests/data/bootstrap_user_data	2011-12-02 14:44:25 +0000
@@ -27,6 +27,7 @@
 
 
     exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession
 
     EOF
 
@@ -49,6 +50,7 @@
 
 
     exec python -m juju.agents.provision --nodaemon --logfile /var/log/juju/provision-agent.log
+    --session-file /var/run/juju/provision-agent.zksession
 
     EOF
 

=== modified file 'juju/providers/orchestra/tests/data/launch_user_data'
--- juju/providers/orchestra/tests/data/launch_user_data	2011-12-02 14:44:24 +0000
+++ juju/providers/orchestra/tests/data/launch_user_data	2011-12-02 14:44:25 +0000
@@ -26,6 +26,7 @@
 
 
     exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession
 
     EOF
 

=== modified file 'juju/unit/lifecycle.py'
--- juju/unit/lifecycle.py	2011-12-02 14:44:24 +0000
+++ juju/unit/lifecycle.py	2011-12-02 14:44:25 +0000
@@ -76,10 +76,6 @@
         watches = []
 
         try:
-            # Verify current state
-            assert not self._running, "Already started"
-
-            # Execute the start hook
             if fire_hooks:
                 yield self._execute_hook("config-changed")
                 yield self._execute_hook("start")
@@ -239,9 +235,9 @@
         """Add and remove unit lifecycles per the service relations Determine.
         """
         # changes relation delta of global zk state with our memory state.
-        new_relations = dict([(service_relation.internal_relation_id,
-                               service_relation) for
-                              service_relation in new_relations])
+        new_relations = dict([
+            (service_relation.internal_relation_id, service_relation)
+            for service_relation in new_relations])
         added = set(new_relations.keys()) - set(self._relations.keys())
         removed = set(self._relations.keys()) - set(new_relations.keys())