// 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)} }