How to store an array of values into a variable - postgresql

Question

I have a function which carries out complex load balancing, and I need to first find out the list of idle servers. After that, I have to iterate over a subset of that list, and finally I have to do a lot of complex things, so I don't want to constantly query the list over and over again. See the below as an example (Note that this is PSUEDO CODE ONLY).
CREATE OR REPLACE FUNCTION load_balance (p_company_id BIGINT, p_num_min_idle_servers BIGINT)
RETURNS VOID
AS $$
DECLARE
v_idle_server_ids BIGINT [];
v_num_idle_servers BIGINT;
v_num_idle_servers_to_retire BIGINT;
BEGIN
PERFORM * FROM server FOR UPDATE;
SELECT
count(server.id)
INTO
v_num_idle_servers
WHERE
server.company_id=p_company_id
AND
server.connection_count=0
AND
server.state = 'up';
v_num_idle_servers_to_retire = v_num_idle_servers - p_num_min_idle_servers;
SELECT
server.id
INTO
v_idle_server_ids
WHERE
server.company_id=p_company_id
AND
server.connection_count=0
AND
server.state = 'up'
ORDER BY
server.id;
FOR i in 1..v_num_idle_servers_to_retire
UPDATE
server
SET
state = 'down'
WHERE
server.id = v_idle_server_ids[i];
Question: I was thinking of getting the list of servers and looping over them one by one. Is this possible in postgres?
Note: I tried putting it all in one single query but it gets very, VERY complicated as there are multiple joins and subqueries. For example, a system but have three applications running on three different servers, where the applications know their load but the servers know their company affiliation, so I would need to join the system to the applications and the applications to the servers

Answer 1

Rule of thumb: if you're looping in SQL there's probably a better way.
You want to set state = 'down' until you have a certain number of idle servers.
We can do this in a single statement. Use a Common Table Expression to query your idle servers and feed that to an update from. If you do this a lot you can turn the CTE into a view.
But we need to limit how many we take down based on how many idle servers there are. We can do that with a limit. But update from doesn't take a limit, so we need a second CTE to limit the results.
Here I've hard coded company_id 1 and p_num_min_idle_servers 2.
with idle_servers as (
select id
from server
where server.company_id=1
and connection_count=0
and state = 'up'
),
idle_servers_to_take_down as (
select id
from idle_servers
-- limit doesn't work in update from
limit (select count(*) - 2 from idle_servers)
)
update server
set state = 'down'
from idle_servers_to_take_down
where server.id = idle_servers_to_take_down.id
This has the advantage of being done in one statement avoiding race conditions without having to lock the whole table.
Try it.