Add strategy parameter, allowing to choose best entry.

slapos/networkcachehelper.py

@@ -129,27 +129,32 @@ def helper_upload_network_cached_from_directory(dir_url, cache_url,
       signature_private_key_file, shacache_cert_file, shacache_key_file,
       shadir_cert_file, shadir_key_file)
   finally:
-    # Clean it
-    os.remove(tarpath)
+    # Always clean it
+    if os.path.exists(tarpath):
+      os.remove(tarpath)
   return result
 
 
 def helper_download_network_cached(dir_url, cache_url,
     signature_certificate_list, 
-    directory_key, wanted_metadata_dict={}, required_key_list=[]):
+    directory_key, wanted_metadata_dict={}, required_key_list=[],
+    strategy=None):
   """
   Downloads from a network cache provider.
   Select from shadir directory_key entry matching (at least)
   wanted_metadata_dict and with all metadata keys in required_key_list defined
   and not null.
 
+  if a "strategy" function is given as parameter, use it to choose the best
+  entry between list of matching entries. Otherwise, choose the first.
+  This strategy function takes a list of entries as parameter, and should
+  return the best entry.
+
   If something fails (providor be offline, or hash_string fail), we ignore
   network cached index.
 
   return (file_descriptor, metadata) if succeeded, False otherwise.
   """
-  # XXX It should be possible to give a function as parameter to specify
-  # strategy to choose best matching entry.
   if not(dir_url and cache_url):
     return False
 
@@ -170,6 +175,7 @@ def helper_download_network_cached(dir_url, cache_url,
     json_entry_list = nc.select_generic(directory_key)
     # For each entry shadir sent, chooses only the entry matching all
     # wanted metadata, and having all wanted keys
+    matching_entry_list = []
     for entry, _ in json_entry_list:
       try:
         tags = json.loads(entry)
@@ -188,11 +194,23 @@ def helper_download_network_cached(dir_url, cache_url,
               break
         if not match:
           continue
-        # Everything match. Downloads corresponding file
-        file_descriptor = nc.download(tags.get('sha512'))
-        break
+        # Everything match. Add it to list of matching entries
+        matching_entry_list.append(tags)
       except Exception:
         pass
+    # If a strategy is defined, call it to determine best entry
+    if strategy:
+      best_entry = strategy(matching_entry_list)
+      if not best_entry:
+        logger.info("Can't find best entry matching strategy, selecting "
+            "random one between acceptable ones.")
+        best_entry = matching_entry_list[0]
+    else:
+      best_entry = matching_entry_list[0]
+
+    # download best entry
+    file_descriptor = nc.download(best_entry.get('sha512'))
+
     if file_descriptor is not None:
       return file_descriptor, tags
     else:
@@ -211,14 +229,15 @@ def helper_download_network_cached(dir_url, cache_url,
 
 def helper_download_network_cached_to_file(dir_url, cache_url,
     signature_certificate_list,
-    directory_key, path, wanted_metadata_dict={}):
+    directory_key, path, wanted_metadata_dict={}, required_key_list=[],
+    strategy=None):
   """
   Download a file from network cache. It is the responsibility of caller method
   to check md5.
   """
   result = helper_download_network_cached(dir_url, cache_url,
       signature_certificate_list,
-      directory_key, wanted_metadata_dict)
+      directory_key, wanted_metadata_dict, required_key_list, strategy)
   if result:
     # XXX check if nc filters signature_certificate_list!
     # Creates a file with content to desired path.
@@ -235,7 +254,8 @@ def helper_download_network_cached_to_file(dir_url, cache_url,
 
 def helper_download_network_cached_to_directory(dir_url, cache_url,
     signature_certificate_list,
-    directory_key, path, wanted_metadata_dict={}):
+    directory_key, path, wanted_metadata_dict={}, required_key_list=[],
+    strategy=None):
   """
   Download a tar file from network cache and untar it to specified path.
   """
@@ -246,7 +266,8 @@ def helper_download_network_cached_to_directory(dir_url, cache_url,
     metadata_dict = helper_download_network_cached_to_file(
       dir_url, cache_url,
       signature_certificate_list,
-      directory_key, tarpath, wanted_metadata_dict)
+      directory_key, tarpath, wanted_metadata_dict, required_key_list,
+      strategy)
     if metadata_dict:
       # Untar it to path
       tar = tarfile.open(tarpath)
@@ -258,5 +279,6 @@ def helper_download_network_cached_to_directory(dir_url, cache_url,
 
   finally:
     # Always clean it
-    os.remove(tarpath)
+    if os.path.exists(tarpath):
+      os.remove(tarpath)
   return metadata_dict