core: Add size information to commit metadata

Add a --generate-sizes option to commit to add size information to the
commit metadata.  This will be used by higher level code which wants
to determine the total size necessary for downloading.

7 files changed, 259 insertions, 10 deletions

diff --git a/Makefile-tests.am b/Makefile-tests.am
index ef757712..b634bc43 100644
--- a/Makefile-tests.am
+++ b/Makefile-tests.am
@@ -80,9 +80,10 @@ testmeta_DATA += test-varint.test
 
 if BUILDOPT_GJS
 insttest_SCRIPTS += tests/test-core.js \
+	tests/test-sizes.js \
 	tests/test-sysroot.js \
 	$(NULL)
-testmeta_DATA += test-core.test test-sysroot.test
+testmeta_DATA += test-core.test test-sizes.test test-sysroot.test
 endif
 
 endif
diff --git a/src/libostree/ostree-repo-commit.c b/src/libostree/ostree-repo-commit.c
index 6a4ceb73..5251840f 100644
--- a/src/libostree/ostree-repo-commit.c
+++ b/src/libostree/ostree-repo-commit.c
@@ -32,6 +32,7 @@
 #include "ostree-repo-file-enumerator.h"
 #include "ostree-checksum-input-stream.h"
 #include "ostree-mutable-tree.h"
+#include "ostree-varint.h"
 
 gboolean
 _ostree_repo_ensure_loose_objdir_at (int             dfd,
@@ -170,6 +171,129 @@ commit_loose_object_trusted (OstreeRepo        *self,
   return ret;
 }
 
+typedef struct
+{
+  gsize unpacked;
+  gsize archived;
+} OstreeContentSizeCacheEntry;
+
+static OstreeContentSizeCacheEntry *
+content_size_cache_entry_new (gsize unpacked,
+                              gsize archived)
+{
+  OstreeContentSizeCacheEntry *entry = g_slice_new0 (OstreeContentSizeCacheEntry);
+
+  entry->unpacked = unpacked;
+  entry->archived = archived;
+
+  return entry;
+}
+
+static void
+content_size_cache_entry_free (gpointer entry)
+{
+  if (entry)
+    g_slice_free (OstreeContentSizeCacheEntry, entry);
+}
+
+static void
+repo_store_size_entry (OstreeRepo       *self,
+                       const gchar      *checksum,
+                       gsize             unpacked,
+                       gsize             archived)
+{
+  if (G_UNLIKELY (self->object_sizes == NULL))
+    self->object_sizes = g_hash_table_new_full (g_str_hash, g_str_equal,
+                                                g_free, content_size_cache_entry_free);
+
+  g_hash_table_replace (self->object_sizes,
+                        g_strdup (checksum),
+                        content_size_cache_entry_new (unpacked, archived));
+}
+
+static int
+compare_ascii_checksums_for_sorting (gconstpointer  a_pp,
+                                     gconstpointer  b_pp)
+{
+  char *a = *((char**)a_pp);
+  char *b = *((char**)b_pp);
+
+  return strcmp (a, b);
+}
+
+/**
+ * Create sizes metadata GVariant and add it to the metadata variant given.
+*/
+static gboolean
+add_size_index_to_metadata (OstreeRepo        *self,
+                            GVariant          *original_metadata,
+                            GVariant         **out_metadata,
+                            GCancellable      *cancellable,
+                            GError           **error)
+{
+  gboolean ret = FALSE;
+  gs_unref_variant_builder GVariantBuilder *builder = NULL;
+    
+  if (original_metadata)
+    {
+      builder = ot_util_variant_builder_from_variant (original_metadata, G_VARIANT_TYPE ("a{sv}"));
+    }
+  else
+    {
+      builder = g_variant_builder_new (G_VARIANT_TYPE ("a{sv}"));
+    }
+
+  if (self->object_sizes &&
+      g_hash_table_size (self->object_sizes) > 0)
+    {
+      GHashTableIter entries = { 0 };
+      gchar *e_checksum = NULL;
+      OstreeContentSizeCacheEntry *e_size = NULL;
+      GVariantBuilder index_builder;
+      guint i;
+      gs_unref_ptrarray GPtrArray *sorted_keys = NULL;
+      
+      g_hash_table_iter_init (&entries, self->object_sizes);
+      g_variant_builder_init (&index_builder,
+                              G_VARIANT_TYPE ("a" _OSTREE_OBJECT_SIZES_ENTRY_SIGNATURE));
+
+      /* Sort the checksums so we can bsearch if desired */
+      sorted_keys = g_ptr_array_new ();
+      while (g_hash_table_iter_next (&entries,
+                                     (gpointer *) &e_checksum,
+                                     (gpointer *) &e_size))
+        g_ptr_array_add (sorted_keys, e_checksum);
+      g_ptr_array_sort (sorted_keys, compare_ascii_checksums_for_sorting);
+
+      for (i = 0; i < sorted_keys->len; i++)
+        {
+          guint8 csum[32];
+          const char *e_checksum = sorted_keys->pdata[i];
+          GString *buffer = g_string_new (NULL);
+
+          ostree_checksum_inplace_to_bytes (e_checksum, csum);
+          g_string_append_len (buffer, (char*)csum, 32);
+
+          e_size = g_hash_table_lookup (self->object_sizes, e_checksum);
+          _ostree_write_varuint64 (buffer, e_size->archived);
+          _ostree_write_varuint64 (buffer, e_size->unpacked);
+
+          g_variant_builder_add (&index_builder, "@ay",
+                                 ot_gvariant_new_bytearray ((guint8*)buffer->str, buffer->len));
+          g_string_free (buffer, TRUE);
+        }
+      
+      g_variant_builder_add (builder, "{sv}", "ostree.sizes",
+                             g_variant_builder_end (&index_builder));
+    }
+    
+  ret = TRUE;
+  *out_metadata = g_variant_builder_end (builder);
+  g_variant_ref_sink (*out_metadata);
+
+  return ret;
+}
+
 static gboolean
 write_object (OstreeRepo         *self,
               OstreeObjectType    objtype,
@@ -198,6 +322,8 @@ write_object (OstreeRepo         *self,
   gboolean temp_file_is_regular;
   gboolean is_symlink = FALSE;
   char loose_objpath[_OSTREE_LOOSE_PATH_MAX];
+  gsize unpacked_size = 0;
+  gboolean indexable = FALSE;
 
   g_return_val_if_fail (self->in_transaction, FALSE);
   
@@ -278,6 +404,9 @@ write_object (OstreeRepo         *self,
           gs_unref_object GConverter *zlib_compressor = NULL;
           gs_unref_object GOutputStream *compressed_out_stream = NULL;
 
+          if (self->generate_sizes)
+            indexable = TRUE;
+
           if (!gs_file_open_in_tmpdir_at (self->tmp_dir_fd, 0644,
                                           &temp_filename, &temp_out,
                                           cancellable, error))
@@ -298,10 +427,10 @@ write_object (OstreeRepo         *self,
               /* Don't close the base; we'll do that later */
               g_filter_output_stream_set_close_base_stream ((GFilterOutputStream*)compressed_out_stream, FALSE);
               
-              if (g_output_stream_splice (compressed_out_stream, file_input, 0,
-                                          cancellable, error) < 0)
+              unpacked_size = g_output_stream_splice (compressed_out_stream, file_input,
+                                                      0, cancellable, error);
+              if (unpacked_size < 0)
                 goto out;
-
             }
         }
       else
@@ -341,6 +470,18 @@ write_object (OstreeRepo         *self,
         }
     }
           
+  if (indexable)
+    {
+      gsize archived_size;
+      gs_unref_object GFileInfo *compressed_info =
+        g_file_query_info (temp_file, G_FILE_ATTRIBUTE_STANDARD_SIZE, 0,
+                           cancellable, error);
+      if (!compressed_info)
+        goto out;
+      archived_size = g_file_info_get_size (compressed_info);
+      repo_store_size_entry (self, actual_checksum, unpacked_size, archived_size);
+    }
+
   if (!_ostree_repo_has_loose_object (self, actual_checksum, objtype,
                                       &have_obj, loose_objpath,
                                       cancellable, error))
@@ -1238,15 +1379,21 @@ ostree_repo_write_commit (OstreeRepo      *self,
   gboolean ret = FALSE;
   gs_free char *ret_commit = NULL;
   gs_unref_variant GVariant *commit = NULL;
+  gs_unref_variant GVariant *new_metadata = NULL;
   gs_free guchar *commit_csum = NULL;
   GDateTime *now = NULL;
   OstreeRepoFile *repo_root = OSTREE_REPO_FILE (root);
 
   g_return_val_if_fail (subject != NULL, FALSE);
 
+  /* Add sizes information to our metadata object */
+  if (!add_size_index_to_metadata (self, metadata, &new_metadata,
+                                   cancellable, error))
+    goto out;
+
   now = g_date_time_new_now_utc ();
   commit = g_variant_new ("(@a{sv}@ay@a(say)sst@ay@ay)",
-                          metadata ? metadata : create_empty_gvariant_dict (),
+                          new_metadata ? new_metadata : create_empty_gvariant_dict (),
                           parent ? ostree_checksum_to_bytes_v (parent) : ot_gvariant_new_bytearray (NULL, 0),
                           g_variant_new_array (G_VARIANT_TYPE ("(say)"), NULL, 0),
                           subject, body ? body : "",
@@ -1526,6 +1673,11 @@ write_directory_to_mtree_internal (OstreeRepo                  *self,
 
   g_debug ("Examining: %s", gs_file_get_path_cached (dir));
 
+  if (modifier && modifier->flags & OSTREE_REPO_COMMIT_MODIFIER_FLAGS_GENERATE_SIZES)
+    {
+      self->generate_sizes = TRUE;
+    }
+
   /* If the directory is already in the repository, we can try to
    * reuse checksums to skip checksumming. */
   if (OSTREE_IS_REPO_FILE (dir) && modifier == NULL)
diff --git a/src/libostree/ostree-repo-private.h b/src/libostree/ostree-repo-private.h
index ddd2c7a5..0a3b0a01 100644
--- a/src/libostree/ostree-repo-private.h
+++ b/src/libostree/ostree-repo-private.h
@@ -24,6 +24,8 @@
 
 G_BEGIN_DECLS
 
+#define _OSTREE_OBJECT_SIZES_ENTRY_SIGNATURE "ay"
+
 /**
  * OstreeRepo:
  *
@@ -57,10 +59,12 @@ struct OstreeRepo {
   gboolean in_transaction;
   GHashTable *loose_object_devino_hash;
   GHashTable *updated_uncompressed_dirs;
+  GHashTable *object_sizes;
 
   GKeyFile *config;
   OstreeRepoMode mode;
   gboolean enable_uncompressed_cache;
+  gboolean generate_sizes;
 
   OstreeRepo *parent_repo;
 };
diff --git a/src/libostree/ostree-repo.c b/src/libostree/ostree-repo.c
index 34bd1e88..e7a26bc4 100644
--- a/src/libostree/ostree-repo.c
+++ b/src/libostree/ostree-repo.c
@@ -116,6 +116,7 @@ ostree_repo_finalize (GObject *object)
   g_clear_pointer (&self->txn_refs, g_hash_table_destroy);
   g_clear_pointer (&self->cached_meta_indexes, (GDestroyNotify) g_ptr_array_unref);
   g_clear_pointer (&self->cached_content_indexes, (GDestroyNotify) g_ptr_array_unref);
+  g_clear_pointer (&self->object_sizes, (GDestroyNotify) g_hash_table_unref);
   g_mutex_clear (&self->cache_lock);
   g_mutex_clear (&self->txn_stats_lock);
 
diff --git a/src/libostree/ostree-repo.h b/src/libostree/ostree-repo.h
index 0aa4a3fd..4f3e9d6c 100644
--- a/src/libostree/ostree-repo.h
+++ b/src/libostree/ostree-repo.h
@@ -281,7 +281,8 @@ typedef OstreeRepoCommitFilterResult (*OstreeRepoCommitFilter) (OstreeRepo    *r
  */
 typedef enum {
   OSTREE_REPO_COMMIT_MODIFIER_FLAGS_NONE = 0,
-  OSTREE_REPO_COMMIT_MODIFIER_FLAGS_SKIP_XATTRS = (1 << 0)
+  OSTREE_REPO_COMMIT_MODIFIER_FLAGS_SKIP_XATTRS = (1 << 0),
+  OSTREE_REPO_COMMIT_MODIFIER_FLAGS_GENERATE_SIZES = (1 << 1)
 } OstreeRepoCommitModifierFlags;
 
 /**
diff --git a/src/ostree/ot-builtin-commit.c b/src/ostree/ot-builtin-commit.c
index e9c030d7..6f632056 100644
--- a/src/ostree/ot-builtin-commit.c
+++ b/src/ostree/ot-builtin-commit.c
@@ -45,6 +45,7 @@ static gboolean opt_table_output;
 static char **opt_key_ids;
 static char *opt_gpg_homedir;
 #endif
+static gboolean opt_generate_sizes;
 
 static GOptionEntry options[] = {
   { "subject", 's', 0, G_OPTION_ARG_STRING, &opt_subject, "One line subject", "subject" },
@@ -65,6 +66,7 @@ static GOptionEntry options[] = {
   { "gpg-sign", 0, 0, G_OPTION_ARG_STRING_ARRAY, &opt_key_ids, "GPG Key ID to sign the commit with", "key-id"},
   { "gpg-homedir", 0, 0, G_OPTION_ARG_STRING, &opt_gpg_homedir, "GPG Homedir to use when looking for keyrings", "homedir"},
 #endif
+  { "generate-sizes", 0, 0, G_OPTION_ARG_NONE, &opt_generate_sizes, "Generate size information along with commit metadata", NULL },
   { NULL }
 };
 
@@ -284,6 +286,7 @@ ostree_builtin_commit (int argc, char **argv, OstreeRepo *repo, GCancellable *ca
   gs_unref_object OstreeMutableTree *mtree = NULL;
   gs_free char *tree_type = NULL;
   gs_unref_hashtable GHashTable *mode_adds = NULL;
+  OstreeRepoCommitModifierFlags flags = 0;
   OstreeRepoCommitModifier *modifier = NULL;
   OstreeRepoTransactionStats stats;
 
@@ -319,12 +322,17 @@ ostree_builtin_commit (int argc, char **argv, OstreeRepo *repo, GCancellable *ca
       goto out;
     }
 
-  if (opt_owner_uid >= 0 || opt_owner_gid >= 0 || opt_statoverride_file != NULL
+  if (opt_no_xattrs)
+    flags |= OSTREE_REPO_COMMIT_MODIFIER_FLAGS_SKIP_XATTRS;
+  if (opt_generate_sizes)
+    flags |= OSTREE_REPO_COMMIT_MODIFIER_FLAGS_GENERATE_SIZES;
+
+  if (flags != 0
+      || opt_owner_uid >= 0
+      || opt_owner_gid >= 0
+      || opt_statoverride_file != NULL
       || opt_no_xattrs)
     {
-      OstreeRepoCommitModifierFlags flags = 0;
-      if (opt_no_xattrs)
-        flags |= OSTREE_REPO_COMMIT_MODIFIER_FLAGS_SKIP_XATTRS;
       modifier = ostree_repo_commit_modifier_new (flags, commit_filter, mode_adds, NULL);
     }
 
diff --git a/tests/test-sizes.js b/tests/test-sizes.js
new file mode 100644
index 00000000..5cf765fc
--- /dev/null
+++ b/tests/test-sizes.js
@@ -0,0 +1,82 @@
+#!/usr/bin/env gjs
+//
+// Copyright (C) 2013 Colin Walters <walters@verbum.org>
+//
+// This library is free software; you can redistribute it and/or
+// modify it under the terms of the GNU Lesser General Public
+// License as published by the Free Software Foundation; either
+// version 2 of the License, or (at your option) any later version.
+//
+// This library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+// Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public
+// License along with this library; if not, write to the
+// Free Software Foundation, Inc., 59 Temple Place - Suite 330,
+// Boston, MA 02111-1307, USA.
+
+const GLib = imports.gi.GLib;
+const Gio = imports.gi.Gio;
+const OSTree = imports.gi.OSTree;
+
+function assertEquals(a, b) {
+    if (a != b)
+	throw new Error("assertion failed " + JSON.stringify(a) + " == " + JSON.stringify(b));
+}
+
+let testDataDir = Gio.File.new_for_path('test-data');
+testDataDir.make_directory(null);
+testDataDir.get_child('some-file').replace_contents("hello world!", null, false, 0, null);
+testDataDir.get_child('another-file').replace_contents("hello world again!", null, false, 0, null);
+
+let repoPath = Gio.File.new_for_path('repo');
+let repo = OSTree.Repo.new(repoPath);
+repo.create(OSTree.RepoMode.ARCHIVE_Z2, null);
+
+repo.open(null);
+
+let commitModifier = OSTree.RepoCommitModifier.new(OSTree.RepoCommitModifierFlags.GENERATE_SIZES, null);
+
+assertEquals(repo.get_mode(), OSTree.RepoMode.ARCHIVE_Z2);
+
+repo.prepare_transaction(null);
+
+let mtree = OSTree.MutableTree.new();
+repo.write_directory_to_mtree(testDataDir, mtree, commitModifier, null);
+let [,dirTree] = repo.write_mtree(mtree, null);
+let [,commit] = repo.write_commit(null, 'Some subject', 'Some body', null, dirTree, null);
+print("commit => " + commit);
+
+repo.commit_transaction(null, null);
+
+// Test the sizes metadata
+let [,commitVariant] = repo.load_variant(OSTree.ObjectType.COMMIT, commit);
+let metadata = commitVariant.get_child_value(0);
+let sizes = metadata.lookup_value('ostree.sizes', GLib.VariantType.new('aay'));
+let nSizes = sizes.n_children();
+assertEquals(nSizes, 2);
+let expectedUncompressedSizes = [12, 18];
+let foundExpectedUncompressedSizes = 0;
+for (let i = 0; i < nSizes; i++) {
+    let sizeEntry = sizes.get_child_value(i).deep_unpack();
+    assertEquals(sizeEntry.length, 34);
+    let compressedSize = sizeEntry[32];
+    let uncompressedSize = sizeEntry[33];
+    print("compressed = " + compressedSize);
+    print("uncompressed = " + uncompressedSize);
+    for (let j = 0; j < expectedUncompressedSizes.length; j++) {
+	let expected = expectedUncompressedSizes[j];
+	if (expected == uncompressedSize) {
+	    print("Matched expected uncompressed size " + expected);
+	    expectedUncompressedSizes.splice(j, 1);
+	    break;
+	}
+    }
+}
+if (expectedUncompressedSizes.length > 0) {
+    throw new Error("Failed to match expectedUncompressedSizes: " + JSON.stringify(expectedUncompressedSizes));
+}
+
+print("test-sizes complete");