%-*- mode: slang; mode: fold; -*-

% This file provides routines that facilitate the creation of
% master-slave PVM processes

% The following public functions are provided:
%
%     pvm_ms_set_message_callback (fun)
%     pvm_ms_set_hosts ()
%     pvm_ms_run_master (pgms)
%     pvm_ms_slave_exit (exit_status)
%     pvm_ms_run_slave (argv)
%     pvm_ms_set_debug (debug)
%     pvm_ms_set_num_processes_per_host (num_processes)
%     pvm_ms_set_host_max_processes (hostname, max_processes)
%     pvm_ms_kill (mpid, spid)
%     pvm_ms_set_slave_spawned_callback (fun)
%     pvm_ms_set_slave_exit_failed_callback (fun)
%     pvm_ms_set_idle_host_callback (fun)
%     pvm_ms_add_new_slave (slave_argv)

require ("pvm");

static variable Debug = 0;

% Message Tags
static variable TASK_COMPLETE_MSG	= 1000;
static variable TASK_EXIT_MSG		= 1001;
static variable HOST_DELETE_MSG		= 1002;
static variable HOST_ADD_MSG		= 1003;
static variable TASK_STDOUT_MSG		= 1004;
static variable TASK_STATUS_MSG		= 1005;
static variable TASK_KILL_MSG		= 1006;

% Slave status flags
static variable SLAVE_NOTRUN	= 1;   %  waiting in queue
static variable SLAVE_RUNNING	= 2;   %  running
static variable SLAVE_EXITED	= 4;   %  exited

static variable Hosts = NULL;
static variable Default_Max_Processes_Per_Host = 1;
static variable Message_Callback_Hook = NULL;
static variable Slave_Spawned_Callback_Hook = NULL;
static variable Slave_Exit_Failed_Callback_Hook = NULL;
static variable Idle_Host_Callback_Hook = NULL;

%{{{ Routines for dealing with List of Slave Processes

static variable Slave_List;
static variable Slave_Index;

public define pvm_ms_add_new_slave (slave_argv)
{
   variable c = struct
     {
	next,
	  slave_tid, status, host, argv,
	  exit_status, stdout, index
     };

   c.argv = slave_argv;
   c.host = NULL;
   c.slave_tid = NULL;
   c.status = SLAVE_NOTRUN;
   c.next = Slave_List;
   c.exit_status = NULL;
   c.stdout = "";
   c.index = Slave_Index;

   Slave_List = c;
   Slave_Index++;
}

static define create_slave_list (pgms)
{
   Slave_List = NULL;
   Slave_Index = 0;

   foreach (pgms)
     {
	variable slave_argv = ();
	pvm_ms_add_new_slave (slave_argv);
     }
}

static define next_slave (status)
{
   variable c = Slave_List;

   while (c != NULL)
     {
	if (c.status & status)
	  return c;

	c = c.next;
     }
   return c;
}

static define find_slave (tid)
{
   variable c = Slave_List;

   while (c != NULL)
     {
	if (c.slave_tid == tid)
	  return c;

	c = c.next;
     }
   return c;
}

%}}}

%{{{ Routines for dealing with the list of Hosts

static variable Host_List = NULL;
static define add_new_host (hostname, ptid)
{
   variable h = struct
     {
	hostname, ptid, nprocesses, max_processes, next
     };
   h.hostname = hostname;
   h.ptid = ptid;
   h.nprocesses = 0;	       %  if negative, then host is not available
   h.max_processes = Default_Max_Processes_Per_Host;
   h.next = Host_List;
   Host_List = h;
}

static define create_host_list ()
{   
   Host_List = NULL;

   if (Hosts == NULL)
     {
	variable s = pvm_config ();
	_for (0, length (s.hi_name)-1, 1)
	  {
	     variable i = ();
	     add_new_host (s.hi_name[i], s.hi_tid[i]);
	  }
	return;
     }
   
   foreach (Hosts)
     {
	variable host = ();
	add_new_host (host, NULL);
     }
}

static define next_host ()
{
   variable h = Host_List;
   variable h_min, nmin;

   nmin = 0;
   h_min = NULL;
   while (h != NULL)
     {
	if ((h.max_processes - h.nprocesses >= nmin) and (h.nprocesses >= 0))
	  {
	     nmin = h.max_processes - h.nprocesses;
	     h_min = h;
	  }
	h = h.next;
     }

   if (h_min == NULL)
     {
	() = fprintf (stderr, "No HOSTs to work with!  Waiting for more...\n");
	%pvm_exit ();
	%exit (1);
     }

   if (nmin == 0)
     return NULL;

   return h_min;
}

% may return Struct_Type or String_Type
static define find_host_by_ptid (ptid)
{
   variable h = Host_List;
   while (h != NULL)
     {
	if (h.ptid == ptid)
	  return h;
	h = h.next;
     }

   % Not in our list.  Try server's list
   h = pvm_config ();
   variable j = where (h.hi_tid == ptid);
   if (length (j) == 0)
     {
	() = fprintf (stderr, "*** WARNING: host ptid=%d not found\n", ptid);
	return NULL;
     }
   return h.hi_name[j[0]];
}

static define delete_host (host)
{
   host.nprocesses = -1;
   () = fprintf (stderr, "PVMD on host %S(%S) appears down\n", host.hostname,
		 host.ptid);
}

%}}}

%{{{ Routines for handling messages

static define handle_task_exit (msgid, tid)
{
   tid = pvm_unpack (Int_Type);
   variable slave = find_slave (tid);

   slave.host.nprocesses -= 1;

   if (slave.status != SLAVE_EXITED)
     {
	() = fprintf (stderr, "*** Task failed on host %s, will retry\n",
		      slave.host.hostname);
	() = fprintf (stderr, "  (stdout): %S\n", slave.stdout);

	slave.status = SLAVE_NOTRUN;

	if (Slave_Exit_Failed_Callback_Hook != NULL)
	  @Slave_Exit_Failed_Callback_Hook (msgid, tid);

	sleep (1);
	return;
     }

   if (Debug)
     () = fprintf (stderr, "*** Task exit signal caught\n");
   slave.status = SLAVE_EXITED;
}

static define handle_task_complete (msgid, tid)
{
   variable exit_status = pvm_unpack (Int_Type);
   variable slave = find_slave (tid);

   slave.status = SLAVE_EXITED;
   slave.exit_status = exit_status;
   if (Debug)
     () = fprintf (stderr, "*** Task %d completed with status %d\n",
		   tid, exit_status);
}

static define handle_host_delete (msgid, tid)
{
   variable ptid = pvm_unpack (Int_Type);
   variable dtid;

   ERROR_BLOCK
     {
	_clear_error ();
	dtid = -1;
     }

   dtid = pvm_tidtohost (ptid);

   if (Debug)
     vmessage ("**DEBUG: handle_host_delete(%d,0x%X) obtained ptid=0x%X, dtid=0x%X",
	       msgid, tid, ptid, dtid);

   variable host = find_host_by_ptid (dtid);
   if (host == NULL)
     return;

   delete_host (host);
}

static define handle_task_string (msgid, tid)
{
   variable str = pvm_unpack (String_Type);
   vmessage ("Message from Task %d: %s", tid, str);
}

static define handle_task_stdout (msgid, tid)
{
   tid = pvm_unpack (Int_Type);
   variable nbytes = pvm_unpack (Int_Type);
   if (nbytes <= 0)
     return;
   variable str = pvm_unpack (Char_Type, nbytes);
   str = array_to_bstring (str);

   variable slave = find_slave (tid);
   slave.stdout = strcat (slave.stdout, str);;

   if (Debug)
     vmessage ("%d stdout:\n %s", tid, str);
}

static define handle_task_status (msgid, tid)
{
   variable c = Slave_List;
   variable str = "";
   while (c != NULL)
     {
	variable slave_tid = c.slave_tid;
	if (slave_tid == NULL) slave_tid = -1;
	variable status = c.status;
	variable host = c.host;
	if (host != NULL)
	  host = host.hostname;

	if (status == SLAVE_NOTRUN)
	  status = "QUEUED";
	else if (status == SLAVE_RUNNING)
	  status = sprintf ("RUNNING on %S", host);
	else
	  status = sprintf ("EXITED (status=%S)", c.exit_status);

	str += sprintf ("TID: 0x%X\nCmd: %s\nStatus: %s\nHost: %S\n",
			slave_tid, strjoin (c.argv, " "), status, host);

	str += sprintf ("Stdout:\n");
	str += c.stdout;
	str += "\n\n";
	c = c.next;
     }
   pvm_psend (tid, msgid, str);
}

static define handle_host_add (msgid, tid)
{
   variable n = pvm_unpack (Int_Type);
   variable hids = pvm_unpack (Int_Type, n);
   foreach (hids)
     {
	variable hid = ();
	variable host = find_host_by_ptid (hid);
	if (host == NULL)
	  continue;

	if (typeof (host) == String_Type)
	  {
	     vmessage ("Adding host %s", host);
	     add_new_host (host, hid);
	  }
	else vmessage ("Adding host %s", host.hostname);

	if (Idle_Host_Callback_Hook != NULL)
	  @Idle_Host_Callback_Hook ();
     }
}

static define handle_task_kill (msgid, tid)
{
   tid = pvm_unpack (Int_Type);
   if (tid == 0)
     {
	variable c = Slave_List;
	while (c != NULL)
	  {
	     if (c.slave_tid != NULL)
	       {
		  pvm_kill (c.slave_tid);
		  () = fprintf (stderr, "Killed 0x%X\n", c.slave_tid);
	       }
	     c = c.next;
	  }
	() = fprintf (stderr, "Task exiting from kill\n");
	pvm_exit ();
	exit (1);
     }

   pvm_kill (tid);
   () = fprintf (stderr, "Killed 0x%X\n", tid);
}

static define wait_for_message ()
{
   variable bufid = pvm_recv (-1, -1);
   variable msgid, tid;

   (,msgid,tid) = pvm_bufinfo(bufid);

   switch (msgid)
%     {
%     case TASK_STRING_MSG:
%	handle_task_string (msgid, tid);
%    }
     {
      case TASK_COMPLETE_MSG:
	handle_task_complete (msgid, tid);
     }
     {
      case TASK_STDOUT_MSG:
	handle_task_stdout (msgid, tid);
     }
     {
      case TASK_EXIT_MSG:
	handle_task_exit (msgid, tid);
     }
     {
      case HOST_DELETE_MSG:
	handle_host_delete (msgid, tid);
     }
     {
      case HOST_ADD_MSG:
	handle_host_add (msgid, tid);
     }
     {
      case TASK_STATUS_MSG:
	handle_task_status (msgid, tid);
     }
     {
      case TASK_KILL_MSG:
	handle_task_kill (msgid, tid);
     }
     {
	if (Message_Callback_Hook != NULL)
	  {
	     if ((@Message_Callback_Hook) (msgid, tid))
	       return;
	  }
	() = fprintf (stderr, "*** WARNING: unexpected message tag: %d\n", msgid);
     }
}

%}}}

static define spawn_slave (slave, host) %{{{
{
   variable argv = slave.argv;
   variable slave_tid = pvm_spawn (argv, PvmTaskHost, host.hostname, 1);

   if (slave_tid < 0)
     {
	if (slave_tid == PvmNoFile)
	  {
	     () = fprintf (stderr, "Unable to run %s on %S: not found\n",
			   argv[0], host.hostname);
	  }
	() = fprintf (stderr, "Unable to run %s on %S, error=%d: not found\n",
		      argv[0], host.hostname, slave_tid);
	delete_host (host);
	return -1;
     }
   pvm_notify (PvmHostDelete, HOST_DELETE_MSG, slave_tid);
   pvm_notify (PvmTaskExit, TASK_EXIT_MSG, slave_tid);

   slave.host = host;
   slave.slave_tid = slave_tid;
   slave.status = SLAVE_RUNNING;
   host.nprocesses += 1;
   host.ptid = pvm_tidtohost (slave_tid);

   if (Slave_Spawned_Callback_Hook != NULL)
     (@Slave_Spawned_Callback_Hook) (slave_tid, host.hostname, argv);

   if (Debug) vmessage ("task %d spawned on %s", slave_tid, host.hostname);
   return 0;
}

%}}}

public define pvm_ms_set_message_callback (fun)
{
   Message_Callback_Hook = fun;
}

public define pvm_ms_set_slave_spawned_callback (fun)
{
   Slave_Spawned_Callback_Hook = fun;
}

public define pvm_ms_set_slave_exit_failed_callback (fun)
{
   Slave_Exit_Failed_Callback_Hook = fun;
}

public define pvm_ms_set_idle_host_callback (fun)
{
   Idle_Host_Callback_Hook = fun;
}

public define pvm_ms_set_hosts ()
{
   if (_NARGS == 0)
     usage ("pvm_ms_set_hosts (hosts, ...)");

   variable hosts = __pop_args (_NARGS);
   Hosts = [__push_args (hosts)];   
}

static define build_master_exit_status (pgms)
{
   variable n = Slave_Index;
   variable s = Struct_Type[n];
   variable c = Slave_List;

   while (c != NULL)
     {
	variable ss = struct
	  {
	     exit_status, stdout, host
	  };
	ss.exit_status = c.exit_status;
	ss.stdout = c.stdout;
	if (c.host != NULL)
	  ss.host = c.host.hostname;
	s[c.index] = ss;
	c = c.next;
     }

   return s;
}

public define pvm_ms_run_master (pgms)
{
   if (Host_List == NULL) 
     create_host_list ();
   
   create_slave_list (pgms);

   variable master_tid = pvm_mytid ();
   vmessage ("my pid = %d", master_tid);
   pvm_export ("SLANG_MODULE_PATH");
   () = pvm_setopt (PvmOutputTid, master_tid);
   () = pvm_setopt (PvmOutputCode, TASK_STDOUT_MSG);

   pvm_notify (PvmHostAdd, HOST_ADD_MSG, -1);

   while (NULL != (next_slave (SLAVE_NOTRUN | SLAVE_RUNNING)))
     {
	forever
	  {
	     variable slave = next_slave (SLAVE_NOTRUN);
	     if (slave == NULL)
	       break;
	     variable host = next_host ();
	     if (host == NULL)
	       break;
	     if (0 == spawn_slave (slave, host))
	       sleep (0.0);
	  }
	wait_for_message ();
     }

   return build_master_exit_status (pgms);
}

public define pvm_ms_slave_exit (exit_status)
{
   variable parent_tid = pvm_parent();
   pvm_psend (parent_tid, TASK_COMPLETE_MSG, exit_status);
   pvm_exit ();
}

public define pvm_ms_run_slave (argv)
{
   variable mytid = pvm_mytid();
   variable pgm = strjoin (argv, " ");

   pvm_sigterm_enable (1);
   variable exit_status = system (sprintf ("(%s) 2>&1", pgm));

   pvm_ms_slave_exit (exit_status);
   exit (exit_status);
}

public define pvm_ms_set_debug (debug)
{
   Debug = debug;
}

public define pvm_ms_set_num_processes_per_host (num_processes)
{
   if (num_processes < 0)
     {
	vmessage ("*** pvm_ms_set_num_processes_per_host: invalid value:  num_processes = %S", num_processes);
	return;
     }
   
   if (Host_List == NULL)
     create_host_list ();
   
   Default_Max_Processes_Per_Host = num_processes;   
   foreach (Host_List) using ("next")
     {
	variable h = ();
	h.max_processes = num_processes;
     }   
}

public define pvm_ms_set_host_max_processes (hostname, max_processes)
{
   if (max_processes < 0)
     {
	vmessage ("*** pvm_ms_set_host_max_processes: invalid value:  max_processes = %S", max_processes);
	return;
     }   
   
   if (Host_List == NULL) 
     create_host_list ();   
   
   foreach (Host_List) using ("next")
     {
	variable h = ();
	if (h.hostname == hostname)
	  {
	     h.max_processes = max_processes;
	     return;
	  }	
     }   
}

public define pvm_ms_kill (mpid, spid)
{
   pvm_psend (mpid, TASK_COMPLETE_MSG, spid);
}

provide ("pvm_ms");
