pbot/applets/compiler_vm/host/bin/vm-server

#!/usr/bin/env perl

# File: vm-server
#
# Purpose: The compiler server manages the guest virtual machine state and
# listens for incoming compile requests. This server can be run on any remote
# machine. There can be multiple servers using different ports on the same machine.
#
# SPDX-FileCopyrightText: 2021 Pragmatic Software <pragma78@gmail.com>
# SPDX-License-Identifier: MIT

use 5.020;

use warnings;
use strict;

use feature qw(signatures);
no warnings qw(experimental::signatures);

use IO::Socket;
use Net::hostent;
use IPC::Shareable;
use Time::HiRes qw/gettimeofday/;
use Encode;

use constant {
    SERVER_PORT     => $ENV{PBOTVM_SERVER}  // 9000,
    HEARTBEAT_PORT  => $ENV{PBOTVM_HEART}   // 5556,
    DOMAIN_NAME     => $ENV{PBOTVM_DOMAIN}  // 'pbot-vm',
    COMPILE_TIMEOUT => $ENV{PBOTVM_TIMEOUT} // 10,
};

sub vm_revert() {
    return if $ENV{PBOTVM_NOREVERT};
    print "Reverting vm...\n";
    system('time virsh snapshot-revert '.DOMAIN_NAME.' 1');
    print "Reverted.\n";
}

sub execute($command) {
    print "execute ($command)\n";

    my $pid = open(my $fh, '-|', split / /, $command);

    if (not defined $pid) {
        print "Couldn't fork: $!\n";
        return (-13, "[Fatal error]");
    }

    my ($ret, $result) = (0, '');

    ($ret, $result) = eval {
        local $SIG{ALRM} = sub { kill 9, $pid; die "Timed-out: $result\n"; };
        alarm(COMPILE_TIMEOUT);

        while (my $line = <$fh>) {
            $result .= $line;
        }

        my $ret = $? >> 8;
        return ($ret, $result);
    };

    close $fh;
    alarm 0;

    if (my $exception = $@) {
        # handle time-out exception
        if ($exception =~ /Timed-out: (.*)/) {
            return (-13, "[Timed-out] $1");
        }

        # propagate unhandled exception
        die $exception;
    }

    return ($ret, $result);
}

sub connect_to_heartbeat() {
    my $heartbeat;
    my $attempts = 15;

    while (!$heartbeat && $attempts > 0) {
        print "Connecting to heartbeat on port ".HEARTBEAT_PORT." ... ";

        $heartbeat = IO::Socket::INET->new (
            PeerAddr => '127.0.0.1',
            PeerPort => HEARTBEAT_PORT,
            Proto    => 'tcp',
            Type     => SOCK_STREAM,
        );

        if (!$heartbeat) {
            print "failed.\n";
            --$attempts;
            print "Trying again in 2 seconds ($attempts attempts remaining) ...\n" if $attempts > 0;
            sleep 2;
        } else {
            print "success!\n";
        }
    }

    return $heartbeat;
}

sub server_listen($port) {
    my $server = IO::Socket::INET->new (
        Proto     => 'tcp',
        LocalPort => $port,
        Listen    => SOMAXCONN,
        ReuseAddr => 1,
        Reuse     => 1,
    );
    die "Can't setup server: $!" unless $server;
    print "Server $0 accepting clients at :$port\n";
    return $server;
}

sub handle_client($client, $heartbeat) {
    my ($timed_out, $killed) = (0, 0);

    eval {
        # give client 5 seconds to send a line
        local $SIG{ALRM} = sub { die "Client I/O timed-out\n"; };
        alarm 5;

        while (my $line = <$client>) {
            $line =~ s/[\r\n]+$//;
            next if $line =~ m/^\s*$/;

            # give client 5 more seconds
            alarm 5;

            print "Read [$line]\n";

            if (time - $heartbeat > 5) {
                print "Lost heartbeat, ignoring compile attempt.\n";
                print $client "Recovering from previous snippet, try again soon.\n";
                last;
            }

            # disable client time-out
            alarm 0;

            my ($ret, $result) = execute("perl vm-exec $line");

            $result =~ s/\s+$//;
            print "Ret: $ret; result: [$result]\n";

            if ($result =~ m/\[Killed\]$/) {
                $killed = 1;
                $ret = -14;
            }

            if ($ret == -13 && $result =~ m/\[Timed-out\]/) {
                $timed_out = 1;
            }

            print $client $result . "\n";
            last;
        }
    };

    # print client time-out exception
    print "$@" if $@;

    alarm 0;
    close $client;

    print "timed out: $timed_out; killed: $killed\n";
    return $timed_out || $killed;
}

use constant SUCCESS => 0;

sub do_heartbeat() {
    tie my $heartbeat, 'IPC::Shareable', { key => 'heartbeat' };
    tie my $running,   'IPC::Shareable', { key => 'running'   };

    my $heartbeat_monitor = connect_to_heartbeat() || die "No heartbeat.\n";

    while ($running and <$heartbeat_monitor>) {
        $heartbeat = time;
    }

    $heartbeat = 0;
    exit;
}

sub do_server() {
    tie my $heartbeat, 'IPC::Shareable', { key => 'heartbeat' };
    tie my $running,   'IPC::Shareable', { key => 'running'   };

    print "Starting PBot VM Server on port " . SERVER_PORT . "\n";
    my $server = eval { server_listen(SERVER_PORT) };

    if ($@) {
        print STDERR $@;
        $running = 0;
        return;
    }

    while ($running and my $client = $server->accept) {
        $client->autoflush(1);
        my $hostinfo = gethostbyaddr($client->peeraddr);

        print '-' x 20, "\n";
        printf "Connect from %s at %s\n", $client->peerhost, scalar localtime;

        if (handle_client($client, $heartbeat) != SUCCESS) {
            vm_revert();
        }
    }

    print "Shutting down server.\n";
}

sub main() {
    tie my $heartbeat, 'IPC::Shareable', { key => 'heartbeat', create => 1, destroy => 1 };
    tie my $running,   'IPC::Shareable', { key => 'running',   create => 1, destroy => 1 };

    $running   = 1;
    $heartbeat = 0;

    my $heartbeat_pid = fork // die "Heartbeat fork failed: $!";

    if ($heartbeat_pid == 0) {
        do_heartbeat();
    } else {
        do_server();
    }

    print "Waiting for heart to stop...\n";
    waitpid($heartbeat_pid, 0);
    print "Heart stopped.\n";
}

main();