distribute workers based on ratios

distribute.ts

interface Service {
  ratio: number;
  name: string;
}

function assignNumWorkers(services: Service[], maxTotalWorkers: number): Map<string, number> {
  // sub 1 for the orchestrator
  const availableWorkers = maxTotalWorkers - 1;

  // calc total ratio sum
  const totalRatio = services.reduce((sum, service) => sum + service.ratio, 0);

  // init result map
  const workersPerService = new Map<string, number>();

  // 1st pass: calc raw numbers and handle floor
  let totalAssigned = 0;

  for (const service of services) {
    // calc raw number of workers based on ratio
    const rawWorkers = (service.ratio / totalRatio) * availableWorkers;
    // floor so it's an int
    const assignedWorkers = Math.floor(rawWorkers);
    workersPerService.set(service.name, assignedWorkers);
    totalAssigned += assignedWorkers;
  }

  // 2nd pass: distribute rest of workers based on decimal parts
  const remainingToDistribute = availableWorkers - totalAssigned;

  if (remainingToDistribute > 0) {
    // calc decimal parts and sort services by them
    const decimalParts = services.map(service => ({
      name: service.name,
      decimal: (service.ratio / totalRatio) * availableWorkers % 1
    }))
      .sort((a, b) => b.decimal - a.decimal);

    // distribute rest of workers to services with highest decimal parts
    for (let i = 0; i < remainingToDistribute; i++) {
      const serviceName = decimalParts[i % services.length].name;
      workersPerService.set(
        serviceName,
        workersPerService.get(serviceName)! + 1
      );
    }
  }

  return workersPerService;
}

function assert(condition: boolean, message: string) {
  if (!condition) {
    throw new Error(`Assertion failed: ${message}`);
  }
}

/** helper function to sum all workers in a distribution */
function getTotalWorkers(distribution: Map<string, number>): number {
  return Array.from(distribution.values()).reduce((sum, count) => sum + count, 0);
}

// tests
function runTests() {
  console.log("Running tests...");

  // 1: basic
  {
    const services = [
      { name: 'A', ratio: 2 },
      { name: 'B', ratio: 1 },
      { name: 'C', ratio: 1 }
    ];
    const result = assignNumWorkers(services, 9);
    assert(result.get('A') === 4, "Service A should get 4 workers (2/4 ratio of 8 workers)");
    assert(result.get('B') === 2, "Service B should get 2 workers (1/4 ratio of 8 workers)");
    assert(result.get('C') === 2, "Service C should get 2 workers (1/4 ratio of 8 workers)");
    assert(getTotalWorkers(result) === 8, "Total workers should be 8 (9 - 1 orchestrator)");
  }

  // 2: equal
  {
    const services = [
      { name: 'A', ratio: 1 },
      { name: 'B', ratio: 1 },
      { name: 'C', ratio: 1 }
    ];
    const result = assignNumWorkers(services, 7);
    assert(result.get('A') === 2, "Service A should get 2 workers (equal ratio)");
    assert(result.get('B') === 2, "Service B should get 2 workers (equal ratio)");
    assert(result.get('C') === 2, "Service C should get 2 workers (equal ratio)");
    assert(getTotalWorkers(result) === 6, "Total workers should be 6 (7 - 1 orchestrator)");
  }

  // 3: small worker pool
  {
    const services = [
      { name: 'A', ratio: 3 },
      { name: 'B', ratio: 2 },
      { name: 'C', ratio: 1 }
    ];
    const result = assignNumWorkers(services, 4);
    assert(getTotalWorkers(result) === 3, "Total workers should be 3 (4 - 1 orchestrator)");
    assert(result.get('A')! >= 1, "Service A should get at least 1 worker");
    assert(result.get('B')! >= 1, "Service B should get at least 1 worker");
    assert(result.get('C')! >= 0, "Service C might get 0 workers due to small pool");
  }

  // 4: single
  {
    const services = [
      { name: 'A', ratio: 1 }
    ];
    const result = assignNumWorkers(services, 5);
    assert(result.get('A') === 4, "Single service should get all workers minus orchestrator");
    assert(getTotalWorkers(result) === 4, "Total workers should be 4 (5 - 1 orchestrator)");
  }

  // 5: large ratios
  {
    const services = [
      { name: 'A', ratio: 100 },
      { name: 'B', ratio: 50 },
      { name: 'C', ratio: 25 }
    ];
    const result = assignNumWorkers(services, 10);
    assert(getTotalWorkers(result) === 9, "Total workers should be 9 (10 - 1 orchestrator)");
    assert(result.get('A')! > result.get('B')!, "Service A should get more workers than B");
    assert(result.get('B')! > result.get('C')!, "Service B should get more workers than C");
  }

  // 6: even number (again)
  {
    const services = [
      { name: 'a', ratio: 2 },
      { name: 'b', ratio: 2 },
      { name: 'c', ratio: 2 },
      { name: 'd', ratio: 2 },
    ]
    const result = assignNumWorkers(services, 16);
    console.log(result)
    assert(result.get('a')! === 4, `shit got ${result.get('a')}`)
  }

  console.log("All tests passed!");
}

// Run the tests
try {
  runTests();
} catch (e) {
  const error = e instanceof Error ? e : new Error("unknown error")
  console.error("Test failed:", error.message);
}