rucio 37.0.0rc1__py3-none-any.whl

rucio/core/replica_sorter.py

@@ -0,0 +1,362 @@
+# Copyright European Organization for Nuclear Research (CERN) since 2012
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+import shutil
+import socket
+import tarfile
+from collections import OrderedDict
+from datetime import datetime, timedelta
+from math import asin, cos, radians, sin, sqrt
+from pathlib import Path
+from tempfile import TemporaryDirectory, TemporaryFile
+from typing import IO, TYPE_CHECKING, Any, Optional, Union
+from urllib.parse import urlparse
+
+import geoip2.database
+import requests
+from dogpile.cache.api import NO_VALUE
+
+from rucio.common import utils
+from rucio.common.cache import MemcacheRegion
+from rucio.common.config import config_get, config_get_bool, config_get_int
+from rucio.common.constants import SORTING_ALGORITHMS
+from rucio.common.exception import InvalidRSEExpression, SortingAlgorithmNotSupported
+from rucio.core.rse_expression_parser import parse_expression
+
+if TYPE_CHECKING:
+    from _typeshed import StrPath
+
+    from rucio.common.types import IPDict, ReplicaDict
+
+REGION = MemcacheRegion(expiration_time=900, function_key_generator=utils.my_key_generator)
+
+# This product uses GeoLite data created by MaxMind,
+# available from <a href="http://www.maxmind.com">http://www.maxmind.com</a>
+GEOIP_DB_EDITION = 'GeoLite2-City'
+
+
+def extract_file_from_tar_gz(
+        archive_file_obj: IO,
+        file_name: str,
+        destination: 'StrPath'
+) -> None:
+    """
+    Extract one file from the archive and put it at the destination
+
+    archive_fileobj is supposed to be at position 0
+    """
+    with TemporaryDirectory(prefix=file_name) as tmp_dir:
+        tmp_dir = Path(tmp_dir)
+        with tarfile.open(fileobj=archive_file_obj, mode='r:gz') as tfile:
+            tfile.extractall(path=tmp_dir)  # noqa: S202
+            for entry in tfile:
+                if entry.name.find(file_name) > -1:
+                    print('Will move %s to %s' % (tmp_dir / entry.name, destination))
+                    shutil.move(tmp_dir / entry.name, destination)
+
+
+def __download_geoip_db(destination: 'StrPath') -> None:
+    edition_id = GEOIP_DB_EDITION
+    download_url = config_get('core', 'geoip_download_url', raise_exception=False, default=None)
+    verify_tls = config_get_bool('core', 'geoip_download_verify_tls', raise_exception=False, default=True)
+    if not download_url:
+        licence_key = config_get('core', 'geoip_licence_key', raise_exception=False, default=None)
+        if not licence_key:
+            raise Exception('Cannot download GeoIP database: licence key not provided')
+        download_url = 'https://download.maxmind.com/app/geoip_download?edition_id=%s&license_key=%s&suffix=tar.gz' % (edition_id, licence_key)
+
+    result = requests.get(download_url, stream=True, verify=verify_tls)
+    if result and result.status_code in [200, ]:
+        with TemporaryFile() as file_obj:
+            for chunk in result.iter_content(8192):
+                file_obj.write(chunk)
+            file_obj.seek(0)
+
+            extract_file_from_tar_gz(archive_file_obj=file_obj, file_name=f'{edition_id}.mmdb', destination=destination)
+    else:
+        raise Exception('Cannot download GeoIP database: %s, Code: %s, Error: %s' % (edition_id,
+                                                                                     result.status_code,
+                                                                                     result.text))
+
+
+def __geoip_db() -> geoip2.database.Reader:
+    db_path = Path(f'/tmp/{GEOIP_DB_EDITION}.mmdb')
+    db_expire_delay = timedelta(days=config_get_int('core', 'geoip_expire_delay', raise_exception=False, default=30))
+
+    must_download = False
+    if not db_path.is_file():
+        print('%s does not exist. Downloading it.' % db_path)
+        must_download = True
+    elif db_expire_delay and datetime.fromtimestamp(db_path.stat().st_mtime) < datetime.now() - db_expire_delay:
+        print('%s is too old. Re-downloading it.' % db_path)
+        must_download = True
+
+    if must_download:
+        __download_geoip_db(destination=db_path)
+
+    return geoip2.database.Reader(str(db_path))
+
+
+def __get_lat_long(
+        se: str,
+        gi: geoip2.database.Reader
+) -> tuple[Optional[float], Optional[float]]:
+    """
+    Get the latitude and longitude on one host using the GeoLite DB
+    :param se : A hostname or IP.
+    :param gi : A Reader object (geoip2 API).
+    """
+    try:
+        ip = socket.getaddrinfo(se, None)[0][4][0]
+        response = gi.city(ip)
+        return response.location.latitude, response.location.longitude
+    except socket.gaierror as error:
+        # Host definitively unknown
+        print(error)
+    return None, None
+
+
+def __get_distance(
+        se1: str,
+        client_location: 'IPDict',
+        ignore_error: bool
+) -> float:
+    """
+    Get the distance between 2 host using the GeoLite DB
+    :param se1 : A first hostname or IP.
+    :param client_location : contains {'ip', 'fqdn', 'site', 'latitude', 'longitude'}
+    :ignore_error: Ignore exception when the GeoLite DB cannot be retrieved
+    """
+    # does not cache ignore_error, str.lower on hostnames/ips is fine
+    canonical_parties = list(map(lambda x: str(x).lower(), [se1, client_location['ip'], client_location.get('latitude', ''), client_location.get('longitude', '')]))
+    canonical_parties.sort()
+    cache_key = f'replica_sorter:__get_distance|site_distance|{"".join(canonical_parties)}'.replace(' ', '.')
+    cache_val = REGION.get(cache_key)
+    if cache_val is NO_VALUE:
+        try:
+            gi = __geoip_db()
+
+            lat1, long1 = __get_lat_long(se1, gi)
+
+            lat2 = client_location.get('latitude')
+            long2 = client_location.get('longitude')
+
+            if not (lat2 and long2) and client_location['ip'] is not None:
+                lat2, long2 = __get_lat_long(client_location['ip'], gi)
+
+            if lat1 and lat2:
+                long1, lat1, long2, lat2 = map(radians, [long1, lat1, long2, lat2])  # type: ignore (lat2 and long2 might be None)
+                dlon = long2 - long1
+                dlat = lat2 - lat1
+                cache_val = 6378 * 2 * asin(sqrt(sin(dlat / 2)**2 + cos(lat1) * cos(lat2) * sin(dlon / 2)**2))
+                REGION.set(cache_key, cache_val)
+                return cache_val
+        except Exception as error:
+            if not ignore_error:
+                raise error
+        # One host is on the Moon
+        cache_val = 360000
+        REGION.set(cache_key, cache_val)
+    return cache_val  # type: ignore (cache_val should not be None)
+
+
+def __download_custom_distance_table() -> None:
+    """
+    Downloads and parses the custom distance table specified by custom_distance_download_url
+    in the config file. Each line of this CSV file should contain a site name, a RSE name,
+    and a numerical distance value. Any additional fields are silently ignored.
+    """
+    db_path = Path('/tmp/rucio_custom_distance_table.csv')
+    db_expire_delay = timedelta(days=config_get_int('core', 'custom_distance_expire_delay', raise_exception=False, default=30))
+
+    # check if need to download the file
+    must_download = False
+    if not db_path.is_file():
+        print('%s does not exist. Downloading it.' % db_path)
+        must_download = True
+    elif db_expire_delay and datetime.fromtimestamp(db_path.stat().st_mtime) < datetime.now() - db_expire_delay:
+        print('%s is too old. Re-downloading it.' % db_path)
+        must_download = True
+
+    if must_download:
+        download_url = config_get('core', 'custom_distance_download_url', raise_exception=False, default=None)
+        if download_url is None:
+            raise Exception('Cannot download custom distance table: no URL provided')
+        result = requests.get(download_url, stream=True, verify=False)
+        if result and result.status_code in [200, ]:
+            with open(db_path, mode='w') as file_obj:
+                file_obj.write(result.text)
+        else:
+            raise Exception('Cannot download custom distance table: %s, Code: %s, Error: %s' % (download_url,
+                                                                                                result.status_code,
+                                                                                                result.text))
+
+    # parse the local file and add its contents to REGION
+    with open(db_path, mode='r') as f:
+        lines = f.readlines()
+        for line in lines:
+            if line.strip() == "":
+                # ignore blank lines
+                continue
+            bits = line.split(",")
+            if len(bits) < 3:
+                raise Exception('Custom distance table must have at least 3 values per line')
+            # ignore additional fields after first 3 (DUNE has some)
+            site = bits[0].strip()
+            rse = bits[1].strip()
+            distance = float(bits[2].strip())
+            if distance < 0.0 or distance > 1.0:
+                raise Exception('Distances in custom distance table must be in range 0-1')
+            cache_key = f'replica_sorter:__get_distance_custom|site_distance|{rse}|{site}'
+            REGION.set(cache_key, distance)
+
+
+def __get_distance_custom(rse: Union[tuple, str], client_location: 'IPDict') -> float:
+    """
+    Return the distance from a client to a RSE by looking up in custom distance table
+    :param rse: RSE name, or tuple containing replica information with RSE name third
+    :param client_location: location dictionary containing {'ip', 'fqdn', 'site', 'latitude', 'longitude'}
+    :returns: numerical distance value
+    """
+    # get RSE name out of tuple if necessary
+    if isinstance(rse, tuple) and len(rse) == 4:
+        rse = rse[2]
+    cache_key = f'replica_sorter:__get_distance_custom|site_distance|{rse}|{client_location["site"]}'
+    cache_val = REGION.get(cache_key)
+    if not isinstance(cache_val, float):
+        # download the table and add all its values to the cache
+        __download_custom_distance_table()
+        cache_val = REGION.get(cache_key)
+        if not isinstance(cache_val, float):
+            # assume maximum distance if not specified in table
+            cache_val = 1.0
+    return cache_val
+
+
+def site_selector(
+        replicas: dict[str, 'ReplicaDict'],
+        site: str,
+        vo: str
+) -> list[str]:
+    """
+    Return a list of replicas located on one site.
+    :param replicas : A dict with RSEs as values and replicas as keys (URIs).
+    :param site : The name of the site
+    :param vo : The vo within which to search for RSEs
+    """
+    result = []
+    try:
+        rses = parse_expression("site=%s" % site, filter_={'vo': vo})
+    except InvalidRSEExpression:
+        return result
+    except Exception:
+        return result
+    rses = [i['rse'] for i in rses]
+    for replica in replicas:
+        if replicas[replica] in rses:
+            result.append(replica)
+    return result
+
+
+def sort_replicas(
+        dictreplica: dict[str, Any],
+        client_location: 'IPDict',
+        selection: Optional[str] = None
+) -> list[str]:
+    """
+    General sorting method for a dictionary of replicas. Returns the List of replicas.
+
+    :param dictreplica: A dict with replicas as keys (URIs).
+    :param client_location: Location dictionary containing {'ip', 'fqdn', 'site', 'latitude', 'longitude'}
+    :param selection: the selected sorting algorithm.
+    :param default: the default sorting algorithm (random, if not defined).
+    :returns: the keys of dictreplica in a sorted list.
+    """
+    if len(dictreplica) == 0:
+        return []
+
+    if not selection:
+        selection = 'geoip'
+    elif selection not in SORTING_ALGORITHMS:
+        raise SortingAlgorithmNotSupported('Sorting algorithm: %s is not supported. Supported protocols: %s' % (selection, SORTING_ALGORITHMS))
+
+    items = [(key, value) for key, value in dictreplica.items()]
+    # safety check, TODO: remove if all dictreplica values are 4-tuple with priority as second item
+    if isinstance(items[0][1], tuple) and len(items[0][1]) == 4:
+        # sort by value[1], which is the priority
+        items.sort(key=lambda item: item[1][1])
+    dictreplica = OrderedDict(items)
+
+    # all sorts must be stable to preserve the priority (the Python standard sorting functions always are stable)
+    if selection == 'geoip':
+        replicas = sort_geoip(dictreplica, client_location, ignore_error=config_get_bool('core', 'geoip_ignore_error', raise_exception=False, default=True))
+    elif selection == 'custom_table':
+        replicas = sort_custom(dictreplica, client_location)
+    elif selection == 'random':
+        replicas = sort_random(dictreplica)
+
+    return replicas
+
+
+def sort_random(dictreplica: dict[str, Any]) -> list[str]:
+    """
+    Return a list of replicas sorted randomly.
+    :param dictreplica: A dict with replicas as keys (URIs).
+    """
+
+    list_replicas = list(dictreplica.keys())
+    random.shuffle(list_replicas)
+    return list_replicas
+
+
+def sort_geoip(
+        dictreplica: dict[str, Any],
+        client_location: 'IPDict',
+        ignore_error: bool = False
+) -> list[str]:
+    """
+    Return a list of replicas sorted by geographical distance to the client IP.
+    :param dictreplica: A dict with replicas as keys (URIs).
+    :param client_location: Location dictionary containing {'ip', 'fqdn', 'site', 'latitude', 'longitude'}
+    :param ignore_error: Ignore exception when the GeoLite DB cannot be retrieved
+    """
+
+    def distance(pfn: str) -> float:
+        url = urlparse(pfn)
+        if url.scheme == 'root':
+            # handle root proxy urls: root://10.0.0.1//root://192.168.1.1:1094//dpm/....
+            sub_url = urlparse(url.path.lstrip('/'))
+            if sub_url.scheme and sub_url.hostname:
+                url = sub_url
+        return __get_distance(url.hostname, client_location, ignore_error)  # type: ignore (hostname might be None)
+
+    return list(sorted(dictreplica, key=distance))
+
+
+def sort_custom(
+        dictreplica: dict[str, Any],
+        client_location: 'IPDict'
+) -> list[str]:
+    """
+    Return a list of replicas sorted according to the custom distance table.
+    :param dictreplica: A dict with replicas as keys (URIs).
+    :param client_location: Location dictionary containing {'ip', 'fqdn', 'site', 'latitude', 'longitude'}
+    :param ignore_error: Ignore exception when the GeoLite DB cannot be retrieved
+    """
+
+    def distance(pfn: str) -> float:
+        return __get_distance_custom(dictreplica[pfn], client_location)
+
+    return list(sorted(dictreplica, key=distance))