From 2849c8d3eb44b186e9eaed46080796d56e9529f2 Mon Sep 17 00:00:00 2001
From: Joshua Harlow <harlowja@gmail.com>
Date: Sat, 20 Jul 2013 13:06:55 -0700
Subject: Also handle custom handlers correctly.

LP: #1203368
---
 cloudinit/handlers/__init__.py |  9 ++++++---
 cloudinit/helpers.py           |  2 ++
 cloudinit/stages.py            | 28 +++++++++++-----------------
 3 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/cloudinit/handlers/__init__.py b/cloudinit/handlers/__init__.py
index 497d68c5..93df5b61 100644
--- a/cloudinit/handlers/__init__.py
+++ b/cloudinit/handlers/__init__.py
@@ -151,10 +151,12 @@ def walker_handle_handler(pdata, _ctype, _filename, payload):
     try:
         mod = fixup_handler(importer.import_module(modname))
         call_begin(mod, pdata['data'], frequency)
-        # Only register and increment
-        # after the above have worked (so we don't if it
-        # fails)
+        # Only register and increment after the above have worked, so we don't
+        # register if it fails starting.
         handlers.register(mod)
+        # Ensure that it gets finalized by marking said module as having been
+        # initialized correctly.
+        handlers.markings[mod].append('initialized')
         pdata['handlercount'] = curcount + 1
     except:
         util.logexc(LOG, "Failed at registering python file: %s (part "
@@ -230,6 +232,7 @@ def walk(msg, callback, data):
         headers['Content-Type'] = ctype
         callback(data, filename, part.get_payload(decode=True), headers)
         partnum = partnum + 1
+    return partnum
 
 
 def fixup_handler(mod, def_freq=PER_INSTANCE):
diff --git a/cloudinit/helpers.py b/cloudinit/helpers.py
index b91c1290..bd37b8a3 100644
--- a/cloudinit/helpers.py
+++ b/cloudinit/helpers.py
@@ -22,6 +22,7 @@
 
 from time import time
 
+import collections
 import contextlib
 import io
 import os
@@ -281,6 +282,7 @@ class ContentHandlers(object):
 
     def __init__(self):
         self.registered = {}
+        self.markings = collections.defaultdict(list)
 
     def __contains__(self, item):
         return self.is_registered(item)
diff --git a/cloudinit/stages.py b/cloudinit/stages.py
index ed995628..43eaca1b 100644
--- a/cloudinit/stages.py
+++ b/cloudinit/stages.py
@@ -383,21 +383,15 @@ class Init(object):
         # Form our cloud interface
         data = self.cloudify()
 
-        # This list contains the modules initialized (so that we only finalize
-        # ones that were actually initialized)
-        inited_handlers = []
-
         def init_handlers():
             # Init the handlers first
-            called = []
             for (_ctype, mod) in c_handlers.iteritems():
-                if mod in called:
+                if 'initialized' in c_handlers.markings[mod]:
                     # Avoid initing the same module twice (if said module
                     # is registered to more than one content-type).
                     continue
                 handlers.call_begin(mod, data, frequency)
-                inited_handlers.append(mod)
-                called.append(mod)
+                c_handlers.markings[mod].append('initialized')
 
         def walk_handlers():
             # Walk the user data
@@ -413,22 +407,22 @@ class Init(object):
                 # names...
                 'handlercount': 0,
             }
-            handlers.walk(user_data_msg, handlers.walker_callback,
-                          data=part_data)
+            return handlers.walk(user_data_msg, handlers.walker_callback,
+                                 data=part_data)
 
         def finalize_handlers():
             # Give callbacks opportunity to finalize
-            called = []
             for (_ctype, mod) in c_handlers.iteritems():
-                if mod in called:
-                    # Avoid finalizing the same module twice (if said module
-                    # is registered to more than one content-type).
-                    continue
-                if mod not in inited_handlers:
+                mod_markings = c_handlers.markings[mod]
+                if 'initialized' not in mod_markings:
                     # Said module was never inited in the first place, so lets
                     # not attempt to finalize those that never got called.
                     continue
-                called.append(mod)
+                if 'finalized' in mod_markings:
+                    # Avoid finalizing the same module twice (if said module
+                    # is registered to more than one content-type).
+                    continue
+                c_handlers.markings[mod].append('finalized')
                 try:
                     handlers.call_end(mod, data, frequency)
                 except:
-- 
cgit v1.2.3