src.dualinventive.com/go/dinet/vendor/github.com/pebbe/zmq4/examples/ppqueue.go

167 lines
4.1 KiB
Go

//
// Paranoid Pirate queue.
//
package main
import (
zmq "github.com/pebbe/zmq4"
"fmt"
"time"
)
const (
HEARTBEAT_LIVENESS = 3 // 3-5 is reasonable
HEARTBEAT_INTERVAL = 1000 * time.Millisecond // msecs
PPP_READY = "\001" // Signals worker is ready
PPP_HEARTBEAT = "\002" // Signals worker heartbeat
)
// Here we define the worker class; a structure and a set of functions that
// as constructor, destructor, and methods on worker objects:
type worker_t struct {
identity string // Identity of worker
id_string string // Printable identity
expire time.Time // Expires at this time
}
// Construct new worker
func s_worker_new(identity string) worker_t {
return worker_t{
identity: identity,
id_string: identity,
expire: time.Now().Add(HEARTBEAT_INTERVAL * HEARTBEAT_LIVENESS),
}
}
// The ready method puts a worker to the end of the ready list:
func s_worker_ready(self worker_t, workers []worker_t) []worker_t {
for i, worker := range workers {
if self.id_string == worker.id_string {
if i == 0 {
workers = workers[1:]
} else if i == len(workers)-1 {
workers = workers[:i]
} else {
workers = append(workers[:i], workers[i+1:]...)
}
break
}
}
return append(workers, self)
}
// The purge method looks for and kills expired workers. We hold workers
// from oldest to most recent, so we stop at the first alive worker:
func s_workers_purge(workers []worker_t) []worker_t {
now := time.Now()
for i, worker := range workers {
if now.Before(worker.expire) {
return workers[i:] // Worker is alive, we're done here
}
}
return workers[0:0]
}
// The main task is a load-balancer with heartbeating on workers so we
// can detect crashed or blocked worker tasks:
func main() {
frontend, _ := zmq.NewSocket(zmq.ROUTER)
backend, _ := zmq.NewSocket(zmq.ROUTER)
defer frontend.Close()
defer backend.Close()
frontend.Bind("tcp://*:5555") // For clients
backend.Bind("tcp://*:5556") // For workers
// List of available workers
workers := make([]worker_t, 0)
// Send out heartbeats at regular intervals
heartbeat_at := time.Tick(HEARTBEAT_INTERVAL)
poller1 := zmq.NewPoller()
poller1.Add(backend, zmq.POLLIN)
poller2 := zmq.NewPoller()
poller2.Add(backend, zmq.POLLIN)
poller2.Add(frontend, zmq.POLLIN)
for {
// Poll frontend only if we have available workers
var sockets []zmq.Polled
var err error
if len(workers) > 0 {
sockets, err = poller2.Poll(HEARTBEAT_INTERVAL)
} else {
sockets, err = poller1.Poll(HEARTBEAT_INTERVAL)
}
if err != nil {
break // Interrupted
}
for _, socket := range sockets {
switch socket.Socket {
case backend:
// Handle worker activity on backend
// Use worker identity for load-balancing
msg, err := backend.RecvMessage(0)
if err != nil {
break // Interrupted
}
// Any sign of life from worker means it's ready
identity, msg := unwrap(msg)
workers = s_worker_ready(s_worker_new(identity), workers)
// Validate control message, or return reply to client
if len(msg) == 1 {
if msg[0] != PPP_READY && msg[0] != PPP_HEARTBEAT {
fmt.Println("E: invalid message from worker", msg)
}
} else {
frontend.SendMessage(msg)
}
case frontend:
// Now get next client request, route to next worker
msg, err := frontend.RecvMessage(0)
if err != nil {
break // Interrupted
}
backend.SendMessage(workers[0].identity, msg)
workers = workers[1:]
}
}
// We handle heartbeating after any socket activity. First we send
// heartbeats to any idle workers if it's time. Then we purge any
// dead workers:
select {
case <-heartbeat_at:
for _, worker := range workers {
backend.SendMessage(worker.identity, PPP_HEARTBEAT)
}
default:
}
workers = s_workers_purge(workers)
}
}
// Pops frame off front of message and returns it as 'head'
// If next frame is empty, pops that empty frame.
// Return remaining frames of message as 'tail'
func unwrap(msg []string) (head string, tail []string) {
head = msg[0]
if len(msg) > 1 && msg[1] == "" {
tail = msg[2:]
} else {
tail = msg[1:]
}
return
}