// Run the Pool component.
func (pool *Pool) Run() {
conn := pythia.DialRetry(pythia.QueueAddr)
defer conn.Close()
pool.conn = conn
log.Println("Connected to queue", pythia.QueueAddr)
// Tokens is a buffered channel to enforce the capacity. Values do not
// matter.
tokens := make(chan bool, pool.Capacity)
for i := 0; i < pool.Capacity; i++ {
tokens <- true
}
pool.abort = make(chan bool, 1)
var wg sync.WaitGroup
conn.Send(pythia.Message{
Message: pythia.RegisterPoolMsg,
Capacity: pool.Capacity,
})
mainloop:
for {
select {
case msg, ok := <-conn.Receive():
if !ok {
break mainloop
}
switch msg.Message {
case pythia.LaunchMsg:
select {
case <-tokens:
wg.Add(1)
go func(msg pythia.Message) {
pool.doJob(msg.Id, msg.Task, msg.Input)
tokens <- true
wg.Done()
}(msg)
default:
log.Print("Job ", msg.Id, ": capacity exceeded.")
log.Println("Capacity exceeded, cannot handle job.")
conn.Send(pythia.Message{
Message: pythia.DoneMsg,
Id: msg.Id,
Status: pythia.Error,
Output: "Pool capacity exceeded",
})
}
default:
log.Println("Ignoring message", msg.Message)
}
case <-pool.quit:
break mainloop
}
}
conn.Close()
pool.abort <- true
wg.Wait()
}
// Handler function for the server.
func handler(rw http.ResponseWriter, req *http.Request) {
log.Println("Client connected: ", req.URL)
if req.Method != "POST" {
rw.WriteHeader(http.StatusMethodNotAllowed)
return
}
// Reading the task request
body, err := ioutil.ReadAll(req.Body)
if err != nil {
rw.WriteHeader(http.StatusInternalServerError)
return
}
var taskReq taskRequest
if err := json.Unmarshal([]byte(body), &taskReq); err != nil {
rw.WriteHeader(http.StatusBadRequest)
return
}
// Connection to the pool and execution of the task
conn := pythia.DialRetry(pythia.QueueAddr)
defer conn.Close()
content, err := ioutil.ReadFile("tasks/" + taskReq.Tid + ".task")
if err != nil {
rw.WriteHeader(422)
return
}
var task pythia.Task
if err := json.Unmarshal([]byte(content), &task); err != nil {
rw.WriteHeader(http.StatusInternalServerError)
return
}
conn.Send(pythia.Message{
Message: pythia.LaunchMsg,
Id: "test",
Task: &task,
Input: taskReq.Response,
})
if msg, ok := <-conn.Receive(); ok {
switch msg.Status {
case "success":
fmt.Fprintf(rw, msg.Output)
}
return
}
rw.WriteHeader(http.StatusInternalServerError)
}
// DialRetry establishes a test connection, retrying indefinitely.
func DialRetry(t *testing.T, addr net.Addr) *Conn {
return &Conn{t, pythia.DialRetry(addr)}
}