offload-dispatcher/main.py

import copy

from flask import Flask, request
from prometheus_api_client import PrometheusConnect
import random

app = Flask(__name__)
app.config.from_envvar('DISPATCHER_SETTINGS')


def get_prom_client():
    return PrometheusConnect(**app.config["PROMETHEUS"])


@app.route("/")
def hello_world():
    return "Target offload dispatcher"


@app.route("/offload_target")
def offload_target():
    expected_size = int(request.args.get("SIZE_HINT", 20*1024*1024*1024))
    targets = copy.deepcopy(app.config["TARGETS"])
    random.shuffle(targets)
    prom = get_prom_client()
    for target in targets:
        app.logger.info(f"Considering target {target['url']}...")
        result = prom.custom_query(query=target["free_space"]["query"])
        if len(result) != 1:
            continue
        result = result[0]
        target_free_space = int(result['value'][1])
        expected_free_space = target_free_space - expected_size
        app.logger.info(f"Available space: {target_free_space} Expected size: {expected_size}")
        if expected_free_space < target["free_space"].get("minimum", 1024 * 1024 * 1024 * 500):
            continue
        return {"url": target["url"]}
    return "Unable to allocate target", 500