/*
 * Copyright (C) 1996-1998 by the Board of Trustees
 *    of Leland Stanford Junior University.
 * 
 * This file is part of the SimOS distribution. 
 * See LICENSE file for terms of the license. 
 *
 */


/* simos_cpualloc -- a simple allocator to round-robin simos jobs across cpus
 *
 * When each simos process starts up it tries to connect to the TCP socket
 * exported by this process.  If it succeeds it passes in its pid and
 * a set of cpus it should not be scheduled on (so the user can avoid non-simos
 * processes) and gets back a cpu.
 *
 * The cpu allocator keeps a cpu->pid mapping.  When it has allocated all cpus
 * it cleans out the mapping table by checking whether those pids are still
 * alive.
 */

#include <sys/types.h>
#include <sys/socket.h>
#include <sys/sysmp.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <signal.h>
#include <stdlib.h>

#define MAXSKIPS 16
typedef struct Simos_cpu_request {
   uint pid;
   uint num_skips;
   uint skips[MAXSKIPS];
} Simos_cpu_request;

#define DEFAULT_PORT 2330

static int alloccpu(Simos_cpu_request*);

/* command line options */

int ncpus = -1;   /* default is to query number of cpus on the current system */
int portnum = DEFAULT_PORT;


void parseopts(int argc, char** argv)
{
   int c;
   extern char* optarg;
   
   while ((c = getopt(argc, argv, "p:c:")) != EOF) {
      switch (c) {
      case 'p':
	 portnum = atoi(optarg);
	 break;
      case 'c':
	 ncpus = atoi(optarg);
	 break;
      default:
	 fprintf(stderr, "usage: %s [-p portnum] [-c numcpus]\n");
	 exit(1);
	 break;
      }
   }

   if (ncpus == -1) {
      ncpus = sysmp(MP_NPROCS);
      if (ncpus <= 0) {
	 perror("simos_cpualloc: MP_NPROCS");
	 exit(1);
      }
   }
}

int main(int argc, char** argv)
{
   int s, tmp;
   struct sockaddr_in addr, from;

   parseopts(argc, argv);

   s = socket(PF_INET, SOCK_STREAM, 0);
   if (s < 0) {
      perror("simos_cpualloc: creating socket");
      exit(1);
   }

   /* Allow rapid reuse of this port. */
   tmp = 1;
   if (setsockopt (s, SOL_SOCKET, SO_REUSEADDR, (char *)&tmp,
		   sizeof(tmp)) < 0) {
      perror("simos_cpualloc: setsockopt");
      exit(1);
   }

   addr.sin_family = PF_INET;
   addr.sin_port = htons(portnum);
   addr.sin_addr.s_addr = INADDR_ANY;

   while (bind (s, &addr, sizeof (addr))
	  || listen (s, 5)) {
      portnum++;
      addr.sin_port = htons(portnum);
   }
   fprintf(stderr, "simos_cpualloc: portnum is %d\n", portnum);

   /* go to background */

   if (_daemonize(0, 2, s, -1) < 0) {
      exit(1);
   }

   while (1) {
      int fromlen = sizeof(from);
      Simos_cpu_request argreq;
      int nbytes, result;
      int fd;
      int len = 0;
      char* buf = (char*) &argreq;

      fd  = accept (s, &from, &fromlen);
      if (fd < 0) {
	 perror("simos_cpualloc: accept");
	 exit(1);
      }

       /* Tell TCP not to delay small packets. */
      tmp = 1;
      setsockopt (fd, IPPROTO_TCP, TCP_NODELAY,
		  (char *)&tmp, sizeof(tmp));

      while (len < sizeof(argreq)) {
	 nbytes = read(fd, buf, sizeof(argreq));

	 if (nbytes < 0) {
	    perror("simos_cpualloc: read");
	    goto end_connection;
	 }
	 len += nbytes;
      }

      result = alloccpu(&argreq);
      if (write(fd, &result, sizeof(result)) != sizeof(result)) {
	 perror("simos_cpualloc: write");
      }

 end_connection:
      close(fd);
   }
   
   return 0;
}

/*****************************************************************/

#define MAXCPUS 64
#define PIDSPERCPU 4

typedef struct cpureq {
   int numpids;
   int skip;
   int pids[PIDSPERCPU];
} cpureq;

cpureq crec[MAXCPUS];

int alloccpu(Simos_cpu_request* rp)
{
   int cpu, i, j, num_skips;

   for (cpu=0; cpu<ncpus; cpu++) {
      crec[cpu].skip = 0;
   }

   num_skips = rp->num_skips;
   if (num_skips > MAXSKIPS) num_skips = MAXSKIPS;
   for (i=0; i<num_skips; i++) {
      if (rp->skips[i] < ncpus) {
	 crec[rp->skips[i]].skip = 1;
      }
   }

   for (cpu = 0; cpu<ncpus; cpu++) {
      if (crec[cpu].numpids == 0 && ! crec[cpu].skip) {
	 crec[cpu].pids[crec[cpu].numpids++] = rp->pid;
	 return cpu;
      }
   }

   /* table is creeping full.  Eliminate any pids which have exited */

   for (cpu = 0; cpu<ncpus; cpu++) {
      for (i = 0; i<crec[cpu].numpids; i++) {
	 if (kill(crec[cpu].pids[i], 0) < 0) {
	    if (errno != ESRCH) {
	       perror("simos_cpualloc: could not verify pid");
	    }
	    for (j=i; j<crec[cpu].numpids-1; j++) {
	       crec[cpu].pids[j] = crec[cpu].pids[j+1];
	    }
	    crec[cpu].numpids -= 1;
	 }
      }
   }

   for (cpu = 0; cpu<ncpus; cpu++) {
      if (crec[cpu].numpids == 0 && ! crec[cpu].skip) {
	 crec[cpu].pids[crec[cpu].numpids++] = rp->pid;
	 return cpu;
      }
   }

   /* no cpus with no entries.  Pick a random one */

   j = rp->pid;
   j = ((j & 0xf) + ((j>>4)&0xf) + ((j>>8)&0xf)) % ncpus;

   for (cpu = 0; cpu<ncpus; cpu++) {
      i = (j + cpu) % ncpus;
      if (crec[i].numpids < PIDSPERCPU && ! crec[i].skip) {
	 crec[i].pids[crec[i].numpids++] = rp->pid;
	 return i;
      }
   }	 

   /* can't even record this pid */

   for (cpu = 0; cpu<ncpus; cpu++) {
      i = (j + cpu) % ncpus;
      if (crec[i].numpids < PIDSPERCPU && ! crec[i].skip) {
	 crec[i].pids[crec[i].numpids++] = rp->pid;
	 return i;
      }
   }	 

   /* it's his own fault for specifying too many skips */
   
   return j;
}