Move to aiohttp for GeoIP requests
GeoIP requests used to run on the default executor thread and would prevent the application from exiting if the GeoIP service would not respond quickly enough. We witnessed an obvious impact during an incident on geoip.ubuntu.com. Move to aiohttp so that the HTTP calls are non blocking. Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>
This commit is contained in:
parent
58df7d6583
commit
ebdf1db636
2
Makefile
2
Makefile
|
@ -31,7 +31,7 @@ aptdeps:
|
|||
python3-requests-unixsocket python3-jsonschema python3-apport \
|
||||
python3-bson xorriso isolinux python3-aiohttp cloud-init ssh-import-id \
|
||||
curl jq build-essential python3-pytest python3-async-timeout \
|
||||
language-selector-common fuseiso python3-pytest-xdist
|
||||
language-selector-common fuseiso python3-pytest-xdist python3-aioresponses
|
||||
|
||||
.PHONY: install_deps
|
||||
install_deps: aptdeps gitdeps
|
||||
|
|
|
@ -14,13 +14,12 @@
|
|||
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
import aiohttp
|
||||
import logging
|
||||
import enum
|
||||
import requests
|
||||
from xml.etree import ElementTree
|
||||
|
||||
from subiquitycore.async_helpers import (
|
||||
run_in_thread,
|
||||
SingleInstanceTask,
|
||||
)
|
||||
|
||||
|
@ -77,12 +76,12 @@ class HTTPGeoIPStrategy(GeoIPStrategy):
|
|||
geoip.ubuntu.com service. """
|
||||
async def get_response(self) -> str:
|
||||
try:
|
||||
response = await run_in_thread(
|
||||
requests.get, "https://geoip.ubuntu.com/lookup")
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get("https://geoip.ubuntu.com/lookup") as response:
|
||||
response.raise_for_status()
|
||||
except requests.exceptions.RequestException as e:
|
||||
return await response.text()
|
||||
except aiohttp.ClientError as e:
|
||||
raise LookupError from e
|
||||
return response.text
|
||||
|
||||
|
||||
class GeoIP:
|
||||
|
|
|
@ -49,29 +49,16 @@ empty_tz = '<Response><TimeZone></TimeZone></Response>'
|
|||
empty_cc = '<Response><CountryCode></CountryCode></Response>'
|
||||
|
||||
|
||||
class MockGeoIPResponse:
|
||||
def __init__(self, text, status_code=200):
|
||||
self.text = text
|
||||
self.status_code = status_code
|
||||
|
||||
def raise_for_status(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
def requests_get_factory(text):
|
||||
def requests_get(*args, **kwargs):
|
||||
return MockGeoIPResponse(text)
|
||||
return requests_get
|
||||
|
||||
|
||||
class TestGeoIP(SubiTestCase):
|
||||
@mock.patch('requests.get', new=requests_get_factory(xml))
|
||||
def setUp(self):
|
||||
strategy = HTTPGeoIPStrategy()
|
||||
self.geoip = GeoIP(make_app(), strategy)
|
||||
|
||||
async def fn():
|
||||
self.assertTrue(await self.geoip.lookup())
|
||||
|
||||
with aioresponses() as mocked:
|
||||
mocked.get("https://geoip.ubuntu.com/lookup", body=xml)
|
||||
run_coro(fn())
|
||||
|
||||
def test_countrycode(self):
|
||||
|
@ -86,37 +73,42 @@ class TestGeoIPBadData(SubiTestCase):
|
|||
strategy = HTTPGeoIPStrategy()
|
||||
self.geoip = GeoIP(make_app(), strategy)
|
||||
|
||||
@mock.patch('requests.get', new=requests_get_factory(partial))
|
||||
def test_partial_reponse(self):
|
||||
async def fn():
|
||||
self.assertFalse(await self.geoip.lookup())
|
||||
with aioresponses() as mocked:
|
||||
mocked.get("https://geoip.ubuntu.com/lookup", body=partial)
|
||||
run_coro(fn())
|
||||
|
||||
@mock.patch('requests.get', new=requests_get_factory(incomplete))
|
||||
def test_incomplete(self):
|
||||
async def fn():
|
||||
self.assertFalse(await self.geoip.lookup())
|
||||
with aioresponses() as mocked:
|
||||
mocked.get("https://geoip.ubuntu.com/lookup", body=incomplete)
|
||||
run_coro(fn())
|
||||
self.assertIsNone(self.geoip.countrycode)
|
||||
self.assertIsNone(self.geoip.timezone)
|
||||
|
||||
@mock.patch('requests.get', new=requests_get_factory(long_cc))
|
||||
def test_long_cc(self):
|
||||
async def fn():
|
||||
self.assertFalse(await self.geoip.lookup())
|
||||
with aioresponses() as mocked:
|
||||
mocked.get("https://geoip.ubuntu.com/lookup", body=long_cc)
|
||||
run_coro(fn())
|
||||
self.assertIsNone(self.geoip.countrycode)
|
||||
|
||||
@mock.patch('requests.get', new=requests_get_factory(empty_cc))
|
||||
def test_empty_cc(self):
|
||||
async def fn():
|
||||
self.assertFalse(await self.geoip.lookup())
|
||||
with aioresponses() as mocked:
|
||||
mocked.get("https://geoip.ubuntu.com/lookup", body=empty_cc)
|
||||
run_coro(fn())
|
||||
self.assertIsNone(self.geoip.countrycode)
|
||||
|
||||
@mock.patch('requests.get', new=requests_get_factory(empty_tz))
|
||||
def test_empty_tz(self):
|
||||
async def fn():
|
||||
self.assertFalse(await self.geoip.lookup())
|
||||
with aioresponses() as mocked:
|
||||
mocked.get("https://geoip.ubuntu.com/lookup", body=empty_tz)
|
||||
run_coro(fn())
|
||||
self.assertIsNone(self.geoip.timezone)
|
||||
|
|
Loading…
Reference in New Issue