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"


prom = get_prom_client()


@app.route("/offload_target")
def offload_target():
    expected_size = int(request.args.get("SIZE_HINT", 20*1024*1024*1024))
    weights = [x.get("weight", 100) for x in app.config["TARGETS"]]
    target = random.choices(app.config["TARGETS"], weights=weights)[0]
    app.logger.info(f"Considering target {target['url']}...")
    result = prom.custom_query(query=target["free_space"]["query"])
    if len(result) != 1:
        return "Unable to allocate target. Can't get space info.", 500
    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):
        return "Unable to allocate target. Not enough space on selected machine.", 507
    return {"url": target["url"]}