88 lines
1.9 KiB
Go
88 lines
1.9 KiB
Go
package workerserver
|
|
|
|
import (
|
|
"bufio"
|
|
"gitlab.com/saburly/kiviscraplib/structures"
|
|
"log"
|
|
"math/rand"
|
|
"net"
|
|
"time"
|
|
)
|
|
|
|
var requests chan structures.Request
|
|
var workers chan structures.WorkerDescription
|
|
|
|
func ServeWorkers(queue chan structures.Request, end chan<- string) {
|
|
|
|
rand.Seed(time.Now().Unix())
|
|
|
|
workers = make(chan structures.WorkerDescription, 50) // TODO: move to env var
|
|
requests = queue
|
|
|
|
listener, err := net.Listen("tcp", "127.0.0.1:1338")
|
|
if err != nil {
|
|
log.Fatal("tcp server listener error:", err)
|
|
end <- "tcp server"
|
|
return
|
|
}
|
|
|
|
go distributeWork()
|
|
|
|
for {
|
|
conn, err := listener.Accept()
|
|
if err != nil {
|
|
log.Fatal("tcp server accept error", err)
|
|
}
|
|
|
|
go handleConnection(conn)
|
|
}
|
|
}
|
|
|
|
func distributeWork() {
|
|
for {
|
|
request := <-requests
|
|
chosenWorker := <-workers
|
|
chosenWorker.Req <- request
|
|
}
|
|
}
|
|
|
|
func handleConnection(conn net.Conn) {
|
|
clientAddr := conn.RemoteAddr().String()
|
|
workerDescription := structures.WorkerDescription{
|
|
Ip: clientAddr,
|
|
Req: make(chan structures.Request),
|
|
}
|
|
|
|
workers <- workerDescription // add new worker to the end of the queue
|
|
for {
|
|
request := <-workerDescription.Req
|
|
|
|
_, err := conn.Write([]byte("URL " + request.Url + "\n"))
|
|
if err != nil {
|
|
log.Println("Cannot send to " + clientAddr)
|
|
conn.Close()
|
|
requests <- request // return job to some other worker
|
|
return // don't return worker to the queue - not healthy
|
|
}
|
|
|
|
bufferBytes, err := bufio.NewReader(conn).ReadBytes('\n')
|
|
if err != nil {
|
|
log.Println("client left..")
|
|
conn.Close()
|
|
requests <- request // return job to some other worker
|
|
return // don't return worker to the queue - not healthy
|
|
}
|
|
|
|
log.Printf("Recieving response from %s\n", clientAddr)
|
|
response := structures.Response{
|
|
Url: request.Url,
|
|
Content: bufferBytes,
|
|
Err: nil,
|
|
}
|
|
|
|
request.Response <- response
|
|
workers <- workerDescription // finished working and healthy - return to queue
|
|
}
|
|
|
|
}
|