Repository: rucio/rucio

Issue:
Stabilize geoip based replica selection
Motivation
----------
Currently the geoip based replica selection in `list_replicas` does not really work. There were multiple issues in the past, both with how the geoip database is managed, but apparently with the selection as well. 

Modification
------------
- Write test cases
- Identify issues
- Fix issues


Discussion:

Resolving patch:
diff --git a/lib/rucio/core/replica_sorter.py b/lib/rucio/core/replica_sorter.py
--- a/lib/rucio/core/replica_sorter.py
+++ b/lib/rucio/core/replica_sorter.py
@@ -20,18 +20,17 @@
 # - Martin Barisits <martin.barisits@cern.ch>, 2021
 # - Radu Carpa <radu.carpa@cern.ch>, 2021-2022
 
-# This product includes GeoLite data created by MaxMind,
-# available from <a href="http://www.maxmind.com">http://www.maxmind.com</a>
-
 from __future__ import print_function, division
 
-import os
 import random
+import shutil
 import socket
 import tarfile
-import time
 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 TYPE_CHECKING
 from urllib.parse import urlparse
 
@@ -50,46 +49,65 @@
 
 REGION = make_region_memcached(expiration_time=1800, 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, file_name, destination):
+    """
+    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)
+            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(directory, filename):
+def __download_geoip_db(destination):
+    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='NOLICENCE')
-        download_url = 'https://download.maxmind.com/app/geoip_download?edition_id=%s&license_key=%s&suffix=tar.gz' % (filename, licence_key)
-    try:
-        os.unlink('%s/%s.tar.gz' % (directory, filename))
-    except OSError:
-        pass
+        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, ]:
-        file_object = open('%s/%s.tar.gz' % (directory, filename), 'wb')
-        for chunk in result.iter_content(8192):
-            file_object.write(chunk)
-        file_object.close()
-        tarfile_name = '%s/%s.tar.gz' % (directory, filename)
-        with tarfile.open(name=tarfile_name, mode='r:gz') as tfile:
-            tfile.extractall(path=directory)
-            for entry in tfile:
-                if entry.name.find('%s.mmdb' % filename) > -1:
-                    print('Will move %s/%s to %s/%s' % (directory, entry.name, directory, entry.name.split('/')[-1]))
-                    os.rename('%s/%s' % (directory, entry.name), '%s/%s' % (directory, entry.name.split('/')[-1]))
+        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' % (filename,
+        raise Exception('Cannot download GeoIP database: %s, Code: %s, Error: %s' % (edition_id,
                                                                                      result.status_code,
                                                                                      result.text))
 
 
-def __get_geoip_db(directory, filename):
-    if directory.endswith('/'):
-        directory = directory[:-1]
-    if not os.path.isfile('%s/%s.mmdb' % (directory, filename)):
-        print('%s does not exist. Downloading it.' % (filename))
-        __download_geoip_db(directory, filename)
-    elif time.time() - os.stat('%s/%s.mmdb' % (directory, filename)).st_atime > 30 * 86400:
-        print('%s is too old. Re-downloading it.' % (filename))
-        __download_geoip_db(directory, filename)
-    return
+def __geoip_db():
+    db_path = Path(f'/tmp/{GEOIP_DB_EDITION}.mmdb')
+    db_expire_delay = timedelta(days=30)
+
+    must_download = False
+    if not db_path.is_file():
+        print('%s does not exist. Downloading it.' % db_path)
+        must_download = True
+    elif 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, gi):
@@ -121,15 +139,11 @@ def __get_distance(se1, client_location, ignore_error):
     cache_key = f'replica_sorter:__get_distance|site_distance|{"".join(canonical_parties)}'
     cache_val = REGION.get(cache_key)
     if cache_val is NO_VALUE:
-        directory = '/tmp'
-        ipv6_filename = 'GeoLite2-City'
         try:
-            __get_geoip_db(directory, ipv6_filename)
-
-            gi = geoip2.database.Reader('%s/%s' % (directory, '%s.mmdb' % ipv6_filename))
+            gi = __geoip_db()
 
             lat1, long1 = __get_lat_long(se1, gi)
-            if client_location['latitude'] and client_location['longitude']:
+            if client_location.get('latitude') and client_location.get('longitude'):
                 lat2 = client_location['latitude']
                 long2 = client_location['longitude']
             else:
diff --git a/lib/rucio/web/rest/flaskapi/v1/replicas.py b/lib/rucio/web/rest/flaskapi/v1/replicas.py
--- a/lib/rucio/web/rest/flaskapi/v1/replicas.py
+++ b/lib/rucio/web/rest/flaskapi/v1/replicas.py
@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-# Copyright 2021 CERN
+# Copyright 2021-2022 CERN
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,7 +17,7 @@
 # - Benedikt Ziemons <benedikt.ziemons@cern.ch>, 2021
 # - Thomas Beermann <thomas.beermann@cern.ch>, 2021
 # - Eric Vaandering <ewv@fnal.gov>, 2021
-# - Radu Carpa <radu.carpa@cern.ch>, 2021
+# - Radu Carpa <radu.carpa@cern.ch>, 2021-2022
 # - Ilija Vukotic <ivukotic@cern.ch>, 2021
 # - Martin Barisits <martin.barisits@cern.ch>, 2021
 # - Joel Dierkes <joel.dierkes@cern.ch>, 2021
@@ -46,18 +46,21 @@
     request_auth_env, response_headers, generate_http_error_flask, ErrorHandlingMethodView, json_parameters, param_get
 
 
-def _sort_pfns_by_priority(pfns, sorted_pfns):
+def _sorted_with_priorities(replicas, sorted_pfns, limit=None):
     """
-    Sets the corresponding priority values for the given list of pfns and yields
-    the index of the element as well as the element.
+    Pick up to "limit" replicas from "replicas" in the order given by sorted_pfns.
+    Sets the corresponding priority in returned replicas.
 
-    :param pfns: List with the pfns which priorities should be set.
+    :param replicas: Dictionary {pfn: replica_definition}.
     :param sorted_pfns: Sorted list of pfns.
+    :param limit: only return this many replicas
     :yields: index and corresponding pfn
     """
     for idx, pfn in enumerate(sorted_pfns, start=1):
-        pfns[pfn]['priority'] = idx
-        yield idx, pfn
+        if limit is None or idx <= limit:
+            replica = replicas[pfn]
+            replica['priority'] = idx
+            yield pfn, replica
 
 
 class Replicas(ErrorHandlingMethodView):
@@ -132,7 +135,7 @@ def generate(vo):
                     replicas = sort_replicas(dictreplica, client_location, selection=select)
 
                     if not metalink:
-                        _ = list(_sort_pfns_by_priority(rfile['pfns'], replicas))
+                        rfile['pfns'] = dict(_sorted_with_priorities(rfile['pfns'], replicas))
                         yield dumps(rfile) + '\n'
                     else:
                         yield ' <file name="' + rfile['name'] + '">\n'
@@ -381,12 +384,12 @@ def generate(request_id, issuer, vo):
                             lanreplicas[pfn] = replica_tuple
                         else:
                             wanreplicas[pfn] = replica_tuple
-                    # Lan replicas sorted by priority; followed by wan replicas sorted by selection criteria
-                    for idx, pfn in _sort_pfns_by_priority(rfile['pfns'],
-                                                           chain(sorted(lanreplicas.keys(), key=lambda pfn: lanreplicas[pfn][1]),
-                                                                 sort_replicas(wanreplicas, client_location, selection=select))):
-                        if limit and limit == idx:
-                            break
+
+                    rfile['pfns'] = dict(_sorted_with_priorities(replicas=rfile['pfns'],
+                                                                 # Lan replicas sorted by priority; followed by wan replicas sorted by selection criteria
+                                                                 sorted_pfns=chain(sorted(lanreplicas.keys(), key=lambda pfn: lanreplicas[pfn][1]),
+                                                                                   sort_replicas(wanreplicas, client_location, selection=select)),
+                                                                 limit=limit))
 
                     if not metalink:
                         yield dumps(rfile, cls=APIEncoder) + '\n'


Test patch:
diff --git a/lib/rucio/tests/inputs/GeoLite2-City-Test.tar.gz b/lib/rucio/tests/inputs/GeoLite2-City-Test.tar.gz
new file mode 100644
Binary files /dev/null and b/lib/rucio/tests/inputs/GeoLite2-City-Test.tar.gz differ
diff --git a/lib/rucio/tests/inputs/__init__.py b/lib/rucio/tests/inputs/__init__.py
new file mode 100644
--- /dev/null
+++ b/lib/rucio/tests/inputs/__init__.py
@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+# Copyright 2022 CERN
+#
+# 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.
+#
+# Authors:
+# - Radu Carpa <radu.carpa@cern.ch>, 2022
+
+import os
+from pathlib import Path
+
+# GeoLite2-City-Test.mmdb downloaded on 26/Jan/2022 from ./test-data/ in
+# https://github.com/maxmind/MaxMind-DB/tree/2f0ef0249245c7f19feffa366793a6fffd529701/
+# Check ./source-data/GeoLite2-City-Test.json in this repository for IPs to use in tests
+GEOIP_LITE2_CITY_TEST_DB = Path(os.path.abspath(__file__)).parent / 'GeoLite2-City-Test.tar.gz'
diff --git a/lib/rucio/tests/test_replica_sorting.py b/lib/rucio/tests/test_replica_sorting.py
--- a/lib/rucio/tests/test_replica_sorting.py
+++ b/lib/rucio/tests/test_replica_sorting.py
@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-# Copyright 2018-2021 CERN
+# Copyright 2018-2022 CERN
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,14 +19,17 @@
 # - Andrew Lister <andrew.lister@stfc.ac.uk>, 2019
 # - Patrick Austin <patrick.austin@stfc.ac.uk>, 2020
 # - Benedikt Ziemons <benedikt.ziemons@cern.ch>, 2020-2021
-# - Radu Carpa <radu.carpa@cern.ch>, 2021
+# - Radu Carpa <radu.carpa@cern.ch>, 2021-2022
 # - Joel Dierkes <joel.dierkes@cern.ch>, 2021
 
 import copy
 import json
+import os
 from unittest import mock
 from urllib.parse import urlparse
+from tempfile import mkstemp
 
+import geoip2.database
 import pytest
 
 from rucio.common.types import InternalAccount, InternalScope
@@ -35,14 +38,67 @@
 from rucio.core.replica import add_replicas, delete_replicas
 from rucio.core.rse import add_rse, del_rse, add_rse_attribute, add_protocol, del_rse_attribute
 from rucio.tests.common import rse_name_generator, headers, auth, vohdr, Mime, accept
+from rucio.tests.inputs import GEOIP_LITE2_CITY_TEST_DB
+
+LOCATION_TO_IP = {
+    'Switzerland': '2a02:d000::1',
+    'Romania': '2a02:e940::1',
+    'Austria': '2a02:da80::1',
+    'United Kingdom': '81.2.69.142',
+    'China': '175.16.199.0',
+    'United States': '216.160.83.56',
+    'Japan': '2001:258::1',
+    'Taiwan': '2001:288::1',
+    'Israel': '2a02:cf80::1',
+    'Finland': '2a02:d200::1',
+    'United Arab Emirates': '2a02:f400::1',
+    'Libya': '2a02:e700::1',
+}
 
 base_rse_info = [
-    {'site': 'APERTURE', 'address': 'aperture.com'},
-    {'site': 'BLACKMESA', 'address': 'blackmesa.com'},
+    {'site': 'APERTURE', 'address': 'aperture.com', 'ip': LOCATION_TO_IP['Austria']},
+    {'site': 'BLACKMESA', 'address': 'blackmesa.com', 'ip': LOCATION_TO_IP['Japan']},
 ]
 schemes = ['root', 'gsiftp', 'davs']
 
 
+@pytest.fixture
+def mock_geoip_db():
+    temp_fd, temp_db_path = mkstemp()
+    os.close(temp_fd)
+    try:
+        with open(GEOIP_LITE2_CITY_TEST_DB, 'rb') as archive_file:
+            replica_sorter.extract_file_from_tar_gz(archive_file_obj=archive_file,
+                                                    file_name='GeoLite2-City-Test.mmdb',
+                                                    destination=temp_db_path)
+        with mock.patch('rucio.core.replica_sorter.__geoip_db', side_effect=lambda: geoip2.database.Reader(temp_db_path)):
+            yield geoip2.database.Reader(temp_db_path)
+    finally:
+        os.unlink(temp_db_path)
+
+
+@pytest.fixture
+def mock_get_lat_long():
+    def _get_lat_long_mock(se, gi):
+        ip = None
+        if se in LOCATION_TO_IP.values():
+            ip = se
+        else:
+            # Try to map the hostname to one of the test RSES fake ip
+            for rse_info in base_rse_info:
+                if rse_info['address'] in se:
+                    ip = rse_info['ip']
+                    break
+        if ip:
+            response = gi.city(ip)
+            return response.location.latitude, response.location.longitude
+
+        raise Exception("Unknown ip provided, fail the test")
+
+    with mock.patch('rucio.core.replica_sorter.__get_lat_long', side_effect=_get_lat_long_mock):
+        yield
+
+
 @pytest.fixture
 def protocols_setup(vo):
     rse_info = copy.deepcopy(base_rse_info)
@@ -65,7 +121,7 @@ def protocols_setup(vo):
         assert len(site_rses) > 0
         assert rse_info[idx]['id'] in [rse['id'] for rse in site_rses]
 
-    add_protocol(rse_info[0]['id'], {'scheme': schemes[0],
+    add_protocol(rse_info[0]['id'], {'scheme': 'root',
                                      'hostname': ('root.%s' % base_rse_info[0]['address']),
                                      'port': 1409,
                                      'prefix': '//test/chamber/',
@@ -73,7 +129,7 @@ def protocols_setup(vo):
                                      'domains': {
                                          'lan': {'read': 1, 'write': 1, 'delete': 1},
                                          'wan': {'read': 1, 'write': 1, 'delete': 1}}})
-    add_protocol(rse_info[0]['id'], {'scheme': schemes[2],
+    add_protocol(rse_info[0]['id'], {'scheme': 'davs',
                                      'hostname': ('davs.%s' % base_rse_info[0]['address']),
                                      'port': 443,
                                      'prefix': '/test/chamber/',
@@ -81,7 +137,7 @@ def protocols_setup(vo):
                                      'domains': {
                                          'lan': {'read': 2, 'write': 2, 'delete': 2},
                                          'wan': {'read': 2, 'write': 2, 'delete': 2}}})
-    add_protocol(rse_info[0]['id'], {'scheme': schemes[1],
+    add_protocol(rse_info[0]['id'], {'scheme': 'gsiftp',
                                      'hostname': ('gsiftp.%s' % base_rse_info[0]['address']),
                                      'port': 8446,
                                      'prefix': '/test/chamber/',
@@ -90,7 +146,7 @@ def protocols_setup(vo):
                                          'lan': {'read': 0, 'write': 0, 'delete': 0},
                                          'wan': {'read': 3, 'write': 3, 'delete': 3}}})
 
-    add_protocol(rse_info[1]['id'], {'scheme': schemes[1],
+    add_protocol(rse_info[1]['id'], {'scheme': 'gsiftp',
                                      'hostname': ('gsiftp.%s' % base_rse_info[1]['address']),
                                      'port': 8446,
                                      'prefix': '/lambda/complex/',
@@ -98,7 +154,7 @@ def protocols_setup(vo):
                                      'domains': {
                                          'lan': {'read': 2, 'write': 2, 'delete': 2},
                                          'wan': {'read': 1, 'write': 1, 'delete': 1}}})
-    add_protocol(rse_info[1]['id'], {'scheme': schemes[2],
+    add_protocol(rse_info[1]['id'], {'scheme': 'davs',
                                      'hostname': ('davs.%s' % base_rse_info[1]['address']),
                                      'port': 443,
                                      'prefix': '/lambda/complex/',
@@ -106,7 +162,7 @@ def protocols_setup(vo):
                                      'domains': {
                                          'lan': {'read': 0, 'write': 0, 'delete': 0},
                                          'wan': {'read': 2, 'write': 2, 'delete': 2}}})
-    add_protocol(rse_info[1]['id'], {'scheme': schemes[0],
+    add_protocol(rse_info[1]['id'], {'scheme': 'root',
                                      'hostname': ('root.%s' % base_rse_info[1]['address']),
                                      'port': 1409,
                                      'prefix': '//lambda/complex/',
@@ -123,6 +179,56 @@ def protocols_setup(vo):
         del_rse(info['id'])
 
 
+@pytest.mark.noparallel(reason='fails when run in parallel, lists replicas and checks for length of returned list')
+@pytest.mark.parametrize("content_type", [Mime.METALINK, Mime.JSON_STREAM])
+def test_sort_geoip_wan_client_location(vo, rest_client, auth_token, protocols_setup, content_type, mock_geoip_db, mock_get_lat_long):
+    """Replicas: test sorting a few WANs via geoip."""
+
+    data = {
+        'dids': [{'scope': f['scope'].external, 'name': f['name'], 'type': 'FILE'} for f in protocols_setup['files']],
+        'schemes': schemes,
+        'sort': 'geoip',
+    }
+
+    first_aut_then_jpn = ['root.aperture.com', 'davs.aperture.com', 'gsiftp.aperture.com', 'gsiftp.blackmesa.com', 'davs.blackmesa.com', 'root.blackmesa.com']
+    first_jpn_then_aut = ['gsiftp.blackmesa.com', 'davs.blackmesa.com', 'root.blackmesa.com', 'root.aperture.com', 'davs.aperture.com', 'gsiftp.aperture.com']
+    for client_location, expected_order in (
+            ('Switzerland', first_aut_then_jpn),
+            ('Romania', first_aut_then_jpn),
+            ('Austria', first_aut_then_jpn),
+            ('United Kingdom', first_aut_then_jpn),
+            ('Libya', first_aut_then_jpn),
+            ('China', first_jpn_then_aut),
+            ('United States', first_jpn_then_aut),
+            ('Japan', first_jpn_then_aut),
+            ('Taiwan', first_jpn_then_aut),
+            ('Israel', first_aut_then_jpn),
+            ('Finland', first_aut_then_jpn),
+            ('United Arab Emirates', first_aut_then_jpn),
+    ):
+        response = rest_client.post(
+            '/replicas/list',
+            headers=headers(auth(auth_token), vohdr(vo), accept(content_type), [('X-Forwarded-For', LOCATION_TO_IP[client_location])]),
+            json=data
+        )
+        assert response.status_code == 200
+        replicas_response = response.get_data(as_text=True)
+        assert replicas_response
+
+        replicas = []
+        pfns = []
+        if content_type == Mime.METALINK:
+            replicas = parse_replicas_from_string(replicas_response)
+            pfns = [s['pfn'] for s in replicas[0]['sources']]
+        elif content_type == Mime.JSON_STREAM:
+            replicas = list(map(json.loads, filter(bool, map(str.strip, replicas_response.splitlines(keepends=False)))))
+            pfns = list(replicas[0]['pfns'])
+
+        print(client_location, pfns)
+        assert len(replicas) == 1
+        assert [urlparse(pfn).hostname for pfn in pfns] == expected_order
+
+
 @pytest.mark.noparallel(reason='fails when run in parallel, lists replicas and checks for length of returned list')
 @pytest.mark.parametrize("content_type", [Mime.METALINK, Mime.JSON_STREAM])
 def test_sort_geoip_wan(vo, rest_client, auth_token, protocols_setup, content_type):
@@ -342,7 +448,7 @@ def fake_get_geoip_db(*args, **kwargs):
     # invalidate cache for __get_distance so that __get_geoip_db is called
     replica_sorter.REGION.invalidate()
 
-    with mock.patch('rucio.core.replica_sorter.__get_geoip_db', side_effect=fake_get_geoip_db) as get_geoip_db_mock:
+    with mock.patch('rucio.core.replica_sorter.__geoip_db', side_effect=fake_get_geoip_db) as get_geoip_db_mock:
         response = rest_client.post(
             '/replicas/list',
             headers=headers(auth(auth_token), vohdr(vo), accept(content_type)),