Be a bit more frugal with memory: serialize XML directly to disk, and,

in the case of snapshot files, serialize incrementally.

svn path=/branches/tk705/; revision=6279

1 files changed, 64 insertions, 49 deletions

diff --git a/rpki/pubdb/models.py b/rpki/pubdb/models.py
index 644f3e3f..b923b3f7 100644
--- a/rpki/pubdb/models.py
+++ b/rpki/pubdb/models.py
@@ -21,12 +21,14 @@ Django ORM models for pubd.
 from __future__ import unicode_literals
 from django.db import models
 from rpki.fields import CertificateField, SundialField
-from lxml.etree import Element, SubElement, tostring as ElementToString
+from lxml.etree import Element, SubElement, ElementTree, xmlfile as XMLFile
 
 import os
 import logging
 import rpki.exceptions
 import rpki.relaxng
+import rpki.x509
+import rpki.POW
 
 logger = logging.getLogger(__name__)
 
@@ -60,6 +62,18 @@ def DERSubElement(elt, name, der, attrib = None, **kwargs):
     return se
 
 
+def sha256_file(f):
+    """
+    Read data from a file-like object, return hex-encoded sha256 hash.
+    """
+
+    h = rpki.POW.Digest(rpki.POW.SHA256_DIGEST)
+    while True:
+        x = f.read(8192)
+        if len(x) == 0:
+            return h.digest().encode("hex")
+        h.update(x)
+
 
 class Client(models.Model):
     client_handle = models.CharField(unique = True, max_length = 255)
@@ -99,7 +113,7 @@ class Session(models.Model):
         delta = Delta(session = self,
                       serial = self.serial + 1,
                       expires = expires)
-        delta.elt = Element(rrdp_tag_delta,
+        delta.xml = Element(rrdp_tag_delta,
                             nsmap = rrdp_nsmap,
                             version = rrdp_version,
                             session_id = self.uuid,
@@ -126,37 +140,32 @@ class Session(models.Model):
 
 
     @staticmethod
-    def _write_rrdp_file(fn, text, rrdp_publication_base, overwrite = False):
-        if overwrite or not os.path.exists(os.path.join(rrdp_publication_base, fn)):
-            tn = os.path.join(rrdp_publication_base, fn + ".%s.tmp" % os.getpid())
-            if not os.path.isdir(os.path.dirname(tn)):
-                os.makedirs(os.path.dirname(tn))
-            with open(tn, "w") as f:
-                f.write(text)
-            os.rename(tn, os.path.join(rrdp_publication_base, fn))
-
-
-    @staticmethod
     def _rrdp_filename_to_uri(fn, rrdp_base_uri):
         return "%s/%s" % (rrdp_base_uri.rstrip("/"), fn)
 
 
-    def _generate_snapshot(self):
-        xml = Element(rrdp_tag_snapshot, nsmap = rrdp_nsmap,
-                      version = rrdp_version,
-                      session_id = self.uuid,
-                      serial = str(self.serial))
-        xml.text = "\n"
-        for obj in self.publishedobject_set.all():
-            DERSubElement(xml, rrdp_tag_publish,
-                          der = obj.der,
-                          uri = obj.uri)
-        rpki.relaxng.rrdp.assertValid(xml)
-        snapshot = ElementToString(xml, pretty_print = True)
-        return snapshot, rpki.x509.sha256(snapshot).encode("hex")
-
-
-    def _generate_update_xml(self, rrdp_base_uri, snapshot_hash):
+    def write_snapshot_file(self, rrdp_publication_base):
+        fn = os.path.join(rrdp_publication_base, self.snapshot_fn)
+        tn = fn + ".%s.tmp" % os.getpid()
+        dn = os.path.dirname(fn)
+        if not os.path.isdir(dn):
+            os.makedirs(dn)
+        with open(tn, "wb+") as f:
+            with XMLFile(f) as xf:
+                with xf.element(rrdp_tag_snapshot, nsmap = rrdp_nsmap,
+                                version = rrdp_version, session_id = self.uuid, serial = str(self.serial)):
+                    xf.write("\n")
+                    for obj in self.publishedobject_set.all():
+                        e = Element(rrdp_tag_publish, nsmap = rrdp_nsmap, uri = obj.uri)
+                        e.text = rpki.x509.base64_with_linebreaks(obj.der)
+                        xf.write(e, pretty_print = True)
+            f.seek(0)
+            h = sha256_file(f)
+        os.rename(tn, fn)
+        return h
+
+
+    def write_notification_xml(self, rrdp_base_uri, snapshot_hash, rrdp_publication_base):
         xml = Element(rrdp_tag_notification, nsmap = rrdp_nsmap,
                       version = rrdp_version,
                       session_id = self.uuid,
@@ -170,29 +179,26 @@ class Session(models.Model):
                        hash =  delta.hash,
                        serial = str(delta.serial))
         rpki.relaxng.rrdp.assertValid(xml)
-        return ElementToString(xml, pretty_print = True)
+        fn = os.path.join(rrdp_publication_base, self.notification_fn)
+        tn = fn + ".%s.tmp" % os.getpid()
+        ElementTree(xml).write(file = tn, pretty_print = True)
+        os.rename(tn, fn)
 
 
-    def synchronize_rrdp_files(self, rrdp_publication_base, rrdp_base_uri, delta):
+    def synchronize_rrdp_files(self, rrdp_publication_base, rrdp_base_uri):
         """
         Write current RRDP files to disk, clean up old files and directories.
         """
 
         current_filenames = self.keep_these_files.copy()
 
-        snapshot_xml, snapshot_hash = self._generate_snapshot()
-        self._write_rrdp_file(self.snapshot_fn, snapshot_xml, rrdp_publication_base)
+        snapshot_hash = self.write_snapshot_file(rrdp_publication_base)
         current_filenames.add(self.snapshot_fn)
 
-        self._write_rrdp_file(delta.fn, delta.xml, rrdp_publication_base)
-        current_filenames.add(delta.fn)
-
         for delta in self.delta_set.all():
             current_filenames.add(delta.fn)
 
-        self._write_rrdp_file(self.notification_fn,
-                              self._generate_update_xml(rrdp_base_uri, snapshot_hash),
-                              rrdp_publication_base, overwrite = True)
+        self.write_notification_xml(rrdp_base_uri, snapshot_hash, rrdp_publication_base),
         current_filenames.add(self.notification_fn)
 
         for root, dirs, files in os.walk(rrdp_publication_base, topdown = False):
@@ -231,10 +237,19 @@ class Delta(models.Model):
         return "%s/deltas/%s.xml" % (self.session.uuid, self.serial)
 
 
-    def activate(self):
-        rpki.relaxng.rrdp.assertValid(self.elt)
-        self.xml = ElementToString(self.elt, pretty_print = True)
-        self.hash = rpki.x509.sha256(self.xml).encode("hex")
+    def activate(self, rrdp_publication_base):
+        rpki.relaxng.rrdp.assertValid(self.xml)
+        fn = os.path.join(rrdp_publication_base, self.fn)
+        tn = fn + ".%s.tmp" % os.getpid()
+        dn = os.path.dirname(fn)
+        if not os.path.isdir(dn):
+            os.makedirs(dn)
+        with open(tn, "wb+") as f:
+            ElementTree(self.xml).write(file = f, pretty_print = True)
+            f.flush()
+            f.seek(0)
+            self.hash = sha256_file(f)
+        os.rename(tn, fn)
         self.save()
         self.session.serial += 1
         self.session.save()
@@ -254,10 +269,10 @@ class Delta(models.Model):
         logger.debug("Publishing %s", uri)
         PublishedObject.objects.create(session = self.session, client = client, der = der, uri = uri,
                                        hash = rpki.x509.sha256(der).encode("hex"))
-        se = DERSubElement(self.elt, rrdp_tag_publish, der = der, uri = uri)
+        se = DERSubElement(self.xml, rrdp_tag_publish, der = der, uri = uri)
         if obj_hash is not None:
             se.set("hash", obj_hash)
-        rpki.relaxng.rrdp.assertValid(self.elt)
+        rpki.relaxng.rrdp.assertValid(self.xml)
 
 
     def withdraw(self, client, uri, obj_hash):
@@ -269,14 +284,14 @@ class Delta(models.Model):
             raise rpki.exceptions.DifferentObjectAtURI("Found different object at %s (old %s, new %s)" % (uri, obj.hash, obj_hash))
         logger.debug("Withdrawing %s", uri)
         obj.delete()
-        SubElement(self.elt, rrdp_tag_withdraw, uri = uri, hash = obj_hash).tail = "\n"
-        rpki.relaxng.rrdp.assertValid(self.elt)
+        SubElement(self.xml, rrdp_tag_withdraw, uri = uri, hash = obj_hash).tail = "\n"
+        rpki.relaxng.rrdp.assertValid(self.xml)
 
 
     def update_rsync_files(self, publication_base):
         from errno import ENOENT
         min_path_len = len(publication_base.rstrip("/"))
-        for pdu in self.elt:
+        for pdu in self.xml:
             assert pdu.tag in (rrdp_tag_publish, rrdp_tag_withdraw)
             fn = self._uri_to_filename(pdu.get("uri"), publication_base)
             if pdu.tag == rrdp_tag_publish:
@@ -301,7 +316,7 @@ class Delta(models.Model):
                         break
                     else:
                         dn = os.path.dirname(dn)
-        del self.elt
+        del self.xml
 
 
 class PublishedObject(models.Model):