Major change in stashcp to allow for methods

Author: djw8605

bin/stashcp

@@ -114,7 +114,10 @@ def doWriteBack(source, destination):
     return curl_exit
 
 
-def doStashCpSingle(sourceFile, destination, debug=False):
+def doStashCpSingle(sourceFile, destination, methods, debug=False):
+    """
+    Perform a single copy from StashCache federation
+    """
 
     global nearest_cache
 
@@ -137,7 +140,56 @@ def doStashCpSingle(sourceFile, destination, debug=False):
 
     # Calculate the starting time
     start1 = int(time.time()*1000)
-    
+
+    # Go through the download methods
+    cur_method = methods[0]
+    success = False
+    for method in methods:
+        cur_method = method
+        if method == "cvmfs":
+            logging.info("Trying CVMFS...")
+            if download_cvmfs(sourceFile, destination, debug, payload):
+                success = True
+                break
+        elif method == "xrootd":
+            logging.info("Trying XrootD...")
+            if download_xrootd(sourceFile, destination, debug, payload):
+                success = True
+                break
+        elif method == "http":
+            logging.info("Trying HTTP...")
+            if download_http(sourceFile, destination, debug, payload):
+                success = True
+                break
+
+    end1 = int(time.time()*1000)
+    payload['start1']=start1
+    payload['end1']=end1
+    payload['timestamp']=end1
+    payload['download_time']=end1-start1
+    if success:
+        payload['status'] = 'Success'
+
+        # Get the final size of the downloaded file
+        if os.path.isdir(destination):
+            destination += "/"
+        dest_dir, dest_filename = os.path.split(destination)
+
+        if dest_filename:
+            final_destination = destination
+        else:
+            final_destination = os.path.join(dest_dir, os.path.basename(sourceFile))
+        payload['filesize'] = os.stat(final_destination).st_size
+        payload['download_size'] = payload['filesize']
+    else:
+        payload['status'] = 'Fail'
+
+    es_send(payload)
+    return success
+
+
+def download_cvmfs(sourceFile, destination, debug, payload):
+
     # First, check if the file is available in CVMFS
     if sourceFile[0] == '/':
         cvmfs_file = os.path.join("/cvmfs/stash.osgstorage.org/", sourceFile[1:])
@@ -149,127 +201,71 @@ def doStashCpSingle(sourceFile, destination, debug=False):
             shutil.copy(cvmfs_file, destination)
             logging.debug("Succesfully copied file from CVMFS!")
             end1 = int(time.time()*1000)
-            dlSz=os.stat(destination).st_size
-            filesize = os.stat(cvmfs_file).st_size
-            dltime=end1-start1
-            destSpace=1
-            status = 'Success'
-            payload['timestamp']=end1
-            payload['host']="CVMFS"
-            payload['download_size']=dlSz
-            payload['filesize'] = filesize
-            payload['download_time']=dltime
-            payload['destination_space']=destSpace
-            payload['status']=status
             payload['tries']=1
-            payload['start1']=start1
-            payload['end1']=end1
             payload['cache']="CVMFS"
-            es_send(payload)
-                
-            return 0
+            payload['host']="CVMFS"
+            return True
 
         except IOError as e:
             logging.error("Unable to copy with CVMFS, even though file exists: %s", str(e))
-    
+            return False
+
     else:
         logging.debug("CVMFS File does not exist")
 
+    return False
+
+def download_xrootd(sourceFile, destination, debug, payload):
+    """
+    Download from the nearest cache, if that fails, fallback to
+    the stash origin.
+    """
+    global nearest_cache
+    global nearest_cache_list
+
+    # Check for xrootd, return quickly if it's not available
     if not check_for_xrootd():
-        return download_with_http(sourceFile, destination, debug)
+        return False
 
     # If the cache is not specified by the command line, then look for the closest
     if not nearest_cache:
         nearest_cache = get_best_stashcache()
+    cache = nearest_cache
     logging.debug("Using Cache %s", nearest_cache)
-
-    # Now check the size of the file with xrootd
-    logging.debug("Checking size of file.")
-    (xrdfs_stdout, xrdfs_stderr) = subprocess.Popen(["xrdfs", main_redirector, "stat", sourceFile], stdout=subprocess.PIPE).communicate()
-    xrdcp_version = subprocess.Popen(['echo $(xrdcp -V 2>&1)'], stdout=subprocess.PIPE, shell=True).communicate()[0][:-1]
-    try:
-        fileSize = int(re.findall(r"Size:   \d+", xrdfs_stdout)[0].split(":   ")[1])
-        logging.debug("Size of the file %s is %s", sourceFile, fileSize)
-        payload['filesize'] = fileSize
-    except (ValueError, IndexError) as e:
-        sys.stderr.write("Unable to find size of file from the origin\n")
-        print str(xrdfs_stdout)
-        sys.stderr.write(str(xrdfs_stderr))
-        sys.stderr.write("\n")
-    
-    payload['xrdcp_version'] = xrdcp_version
-    
-    end1=int(time.time()*1000)
-    payload['end1']=end1
-    payload['start1']=start1
-    
-    start2 = int(time.time()*1000)
 
-    xrd_exit=timed_transfer(filename=sourceFile, debug=debug, destination=destination)
-    
-    end2=int(time.time()*1000)
-
-    dlSz=0
-    if os.path.exists(destination):
-        dlSz=os.stat(destination).st_size
-    destSpace=1
+    xrd_exit = timed_transfer(filename=sourceFile, debug=debug, cache=cache, destination=destination)
 
     payload['xrdexit1']=xrd_exit
-    payload['start2']=start2
-    payload['end2']=end2
-    
 
     if xrd_exit=='0': #worked first try
         logging.debug("Transfer success using %s", nearest_cache)
-        dltime=end2-start2
-        status = 'Success'
-        tries=2
-
-        payload['download_size']=dlSz
-        payload['download_time']=dltime
-        payload['sitename']=sitename
-        payload['destination_space']=destSpace
-        payload['status']=status
-        payload['tries']=tries
-        payload['cache']=nearest_cache
-        if 'filesize' not in payload:
-            payload['filesize']=dlSz
-        es_send(payload)
+        payload['tries'] = 1
+        payload['cache'] = cache
 
     else: #pull from origin
         logging.warning("XrdCP from cache failed on %s, pulling from main redirector", nearest_cache)
-        nearest_cache=main_redirector
-        start3 = int(time.time()*1000)
-        xrd_exit=timed_transfer(filename=sourceFile, debug=debug, destination=destination)
-        end3=int(time.time()*1000)
-        if os.path.exists(destination):
-            dlSz=os.stat(destination).st_size
-        dltime=end3-start3
+        cache = main_redirector
+        xrd_exit=timed_transfer(filename=sourceFile, cache=cache, debug=debug, destination=destination)
+
         if xrd_exit=='0':
             logging.info("Trunk Success")
             status = 'Trunk Sucess'
-            tries=3
+            tries=2
         else:
-            logging.error("stashcp failed after 3 attempts")
+            logging.info("stashcp failed after 2 xrootd attempts")
             status = 'Timeout'
-            tries = 3
-        payload['download_size']=dlSz
-        payload['download_time']=dltime
-        payload['destination_space']=destSpace
+            tries = 2
+
         payload['status']=status
         payload['xrdexit2']=xrd_exit
         payload['tries']=tries
-        payload['start3']=start3
-        payload['end3']=end3
-        payload['cache']=nearest_cache
-        if 'filesize' not in payload:
-            payload['filesize']=dlSz
-        es_send(payload)
+        payload['cache'] = cache
+
         if xrd_exit == '0':
-            return 0
+            return True
         else:
-            return 1
-    return 0
+            return False
+    return True
 
 def check_for_xrootd():
     """
@@ -288,19 +284,17 @@ def check_for_xrootd():
         return False
 
 
-def download_with_http(source, destination, debug):
+def download_http(source, destination, debug, payload):
     """
     Download from the nearest cache with HTTP
     """
     global nearest_cache
     global nearest_cache_list
-    logging.debug("Downloading with HTTP")
 
-    payload = {}
-    payload['filename'] = source
-    sitename = os.environ.setdefault("OSG_SITE_NAME", "siteNotFound")
-    payload['sitename'] = sitename
-    payload.update(parse_job_ad())
+    if not nearest_cache:
+        nearest_cache = get_best_stashcache()
+
+    logging.debug("Downloading with HTTP")
 
     if not nearest_cache:
         nearest_cache = get_best_stashcache()
@@ -359,24 +353,14 @@ def download_with_http(source, destination, debug):
         status = 'Success'
         payload['download_size']=dlSz
         payload['filesize'] = filesize
-    else:
-        status = 'Failure'
-    dltime=end-start
-    destSpace=1
-    payload['timestamp']=end
+
     payload['host']=tried_cache
-    payload['download_time']=dltime
-    payload['destination_space']=destSpace
-    payload['status']=status
     payload['tries']=1
-    payload['start1']=start
-    payload['end1']=end
     payload['cache']=tried_cache
-    es_send(payload)
     if success:
-        return 0
+        return True
     else:
-        return 1
+        return False
 
 
 def parse_job_ad():
@@ -401,7 +385,7 @@ def parse_job_ad():
 
     return temp_list
 
-def dostashcpdirectory(sourceDir, destination, debug=False):
+def dostashcpdirectory(sourceDir, destination, methods, debug=False):
     sourceItems = subprocess.Popen(["xrdfs", stash_origin, "ls", sourceDir], stdout=subprocess.PIPE).communicate()[0].split()
 
     for remote_file in sourceItems:
@@ -410,7 +394,7 @@ def dostashcpdirectory(sourceDir, destination, debug=False):
         if isdir!='0':
             result = dostashcpdirectory(remote_file, destination, debug)
         else:
-            result = doStashCpSingle(remote_file, destination, debug)
+            result = doStashCpSingle(remote_file, destination, methods, debug)
         # Stop transfers if something fails
         if result != 0:
             return result
@@ -420,7 +404,7 @@ def es_send(payload):
 
     # Calculate the curernt timestamp
     payload['timestamp'] = int(time.time()*1000)
-    payload['host'] = payload['cache']
+    #payload['host'] = payload['cache']
 
     def _es_send(payload):
         data = payload
@@ -433,15 +417,15 @@ def es_send(payload):
             f.close()
         except urllib2.URLError, e:
             logging.warning("Error posting to ES: %s", str(e))
-    
+    print payload
     p = multiprocessing.Process(target=_es_send, name="_es_send", args=(payload,))
     p.start()
     p.join(5)
     p.terminate()
 
 
 
-def timed_transfer(filename, destination, debug=False):
+def timed_transfer(filename, destination, cache, debug=False, ):
     """
     Transfer the filename from the cache to the destination using xrdcp
     """
@@ -455,7 +439,7 @@ def timed_transfer(filename, destination, debug=False):
     os.environ.setdefault("XRD_CONNECTIONRETRY", "2")   # How many time should we retry the TCP connection
     os.environ.setdefault("XRD_STREAMTIMEOUT", "30")    # How long to wait for TCP activity
 
-    filepath=nearest_cache+":1094//"+ filename
+    filepath=cache+":1094//"+ filename
     if debug:
         command="xrdcp -d 2 --nopbar -f " + filepath + " " + destination
     else:
@@ -602,6 +586,7 @@ def main():
     parser.add_option('-c', '--cache', dest='cache', help="Cache to use")
     parser.add_option('-j', '--caches-json', dest='caches_json', help="The JSON file containing the list of caches",
                       default=None)
+    parser.add_option('--methods', dest='methods', help="Comma separated list of methods to try, in order.  Default: cvmfs,xrootd,http", default="cvmfs,xrootd,http")
     args,opts=parser.parse_args()
 
     logging.basicConfig(format='%(asctime)s %(name)-12s %(levelname)-8s %(message)s',
@@ -630,12 +615,15 @@ def main():
         nearest_cache = args.cache
         nearest_cache_list = [ args.cache ]
 
+    # Convert the methods
+    methods = args.methods.split(',')
+
     if not args.recursive:
-        result = doStashCpSingle(sourceFile=source, destination=destination, debug=args.debug)
+        result = doStashCpSingle(sourceFile=source, destination=destination, methods = methods, debug=args.debug)
     else:
-        result = dostashcpdirectory(sourceDir = source, destination = destination, debug=args.debug)
+        result = dostashcpdirectory(sourceDir = source, destination = destination, methods = methods, debug=args.debug)
     # Exit with failure
-    sys.exit(result)
+    sys.exit(0 if result else 1)
 
 
 if __name__ == "__main__":