<commit_before>package dendrite

import (
	zmq "github.com/pebbe/zmq4"
	"log"
	"sync"
	"time"
)

type controlType int

const (
	workerShutdownReq controlType = iota
	workerShutdownAllowed
	workerShutdownDenied
	workerShutdownConfirm
	workerRegisterReq
	workerRegisterAllowed
	workerRegisterDenied
)

// ZeroMQ Transport implementation
type ZMQTransport struct {
	lock              *sync.Mutex
	minHandlers       int
	maxHandlers       int
	incrHandlers      int
	activeRequests    int
	control_c         chan *workerComm
	dealer_sock       *zmq.Socket
	router_sock       *zmq.Socket
	zmq_context       *zmq.Context
	workerIdleTimeout time.Duration
}

// Creates ZeroMQ transport
// Multiplexes incoming connections which are then processed in separate go routines (workers)
// Multiplexer spawns go routines as needed, but 10 worker routines are created on start
// Every request times out after provided timeout duration
func InitZMQTransport(hostname string, timeout time.Duration) (Transport, error) {
	// initialize ZMQ Context
	context, err := zmq.NewContext()
	if err != nil {
		return nil, err
	}

	// setup router and bind() to tcp address for clients to connect to
	router_sock, err := context.NewSocket(zmq.ROUTER)
	if err != nil {
		return nil, err
	}
	err = router_sock.Bind("tcp://" + hostname)
	if err != nil {
		return nil, err
	}

	// setup dealer for
	dealer_sock, err := context.NewSocket(zmq.DEALER)
	if err != nil {
		return nil, err
	}
	err = dealer_sock.Bind("inproc://dendrite-zmqdealer")
	if err != nil {
		return nil, err
	}
	poller := zmq.NewPoller()
	poller.Add(router_sock, zmq.POLLIN)
	poller.Add(dealer_sock, zmq.POLLIN)

	transport := &ZMQTransport{
		lock:              new(sync.Mutex),
		minHandlers:       10,
		maxHandlers:       1024,
		incrHandlers:      10,
		activeRequests:    0,
		workerIdleTimeout: 10 * time.Second,
		control_c:         make(chan *workerComm),
		dealer_sock:       dealer_sock,
		router_sock:       router_sock,
		zmq_context:       context,
	}
	// proxy messages between router and dealer
	go func() {
		for {
			sockets, _ := poller.Poll(-1)
			for _, socket := range sockets {
				switch s := socket.Socket; s {
				case router_sock:
					msg, err := s.RecvBytes(0)
					if err != nil {
						log.Println("ERR: TransportListener read bytes on router failed", err)
						continue
					}
					// forward to dealer
					_, err = dealer_sock.SendBytes(msg, 0)
					if err != nil {
						log.Println("ERR: TransportListener forward to dealer failed", err)
						continue
					}
					transport.activeRequests += 1
				case dealer_sock:
					msg, err := s.RecvBytes(0)
					if err != nil {
						log.Println("ERR: TransportListener read bytes on dealer failed", err)
						continue
					}
					// forward up to router
					_, err = router_sock.SendBytes(msg, 0)
					if err != nil {
						log.Println("ERR: TransportListener forward to router failed", err)
						continue
					}
					transport.activeRequests -= 1
				}
			}
		}
	}()

	// Scheduler goroutine keeps track of running workers
	// It spawns new ones if needed, and cancels ones that are idling
	go func() {
		sched_ticker := time.NewTicker(60 * time.Second)
		workers := make(map[*workerComm]bool)
		// fire up initial set of workers
		//for i := 0; i < transport.minHandlers; i++ {
		for i := 0; i < 15; i++ {
			go handleZMQreq(transport)
		}
		for {
			select {
			case comm := <-transport.control_c:
				// worker sent something...
				msg := <-comm.worker_out
				switch {
				case msg == workerRegisterReq:
					if len(workers) == transport.maxHandlers {
						comm.worker_in <- workerRegisterDenied
						log.Println("ERR: TransportListener - max number of workers reached")
						continue
					}
					if _, ok := workers[comm]; ok {
						// worker already registered
						continue
					}
					comm.worker_in <- workerRegisterAllowed
					workers[comm] = true
					log.Println("INFO: TransportListener - registered new worker, total:", len(workers))
				case msg == workerShutdownReq:
					//log.Println("Got shutdown req")
					if len(workers) > transport.minHandlers {
						comm.worker_in <- workerShutdownAllowed
						for _ = range comm.worker_out {
							// wait until worker closes the channel
						}
						delete(workers, comm)
					} else {
						comm.worker_in <- workerShutdownDenied
					}
				}
			case <-sched_ticker.C:
			}
		}
	}()
	return transport, nil
}

type workerComm struct {
	worker_in  chan controlType // worker's input channel for two way communication with scheduler
	worker_out chan controlType // worker's output channel for two way communication with scheduler
	worker_ctl chan controlType // worker's control channel for communication with scheduler
}

func handleZMQreq(transport *ZMQTransport) {
	// setup REP socket
	rep_sock, err := transport.zmq_context.NewSocket(zmq.REP)
	if err != nil {
		log.Println("ERR: TransportListener worker failed to create REP socket", err)
		return
	}
	err = rep_sock.Connect("inproc://dendrite-zmqdealer")
	if err != nil {
		log.Println("ERR: TransportListener worker failed to connect to dealer", err)
		return
	}<commit_msg>adds transport_zmq initializer<commit_after>package dendrite

import (
	zmq "github.com/pebbe/zmq4"
	"log"
	"sync"
	"time"
)

type controlType int

const (
	workerShutdownReq controlType = iota
	workerShutdownAllowed
	workerShutdownDenied
	workerShutdownConfirm
	workerRegisterReq
	workerRegisterAllowed
	workerRegisterDenied
)

// ZeroMQ Transport implementation
type ZMQTransport struct {
	lock              *sync.Mutex
	minHandlers       int
	maxHandlers       int
	incrHandlers      int
	activeRequests    int
	control_c         chan *workerComm
	dealer_sock       *zmq.Socket
	router_sock       *zmq.Socket
	zmq_context       *zmq.Context
	workerIdleTimeout time.Duration
}

// Creates ZeroMQ transport
// Multiplexes incoming connections which are then processed in separate go routines (workers)
// Multiplexer spawns go routines as needed, but 10 worker routines are created on start
// Every request times out after provided timeout duration
func InitZMQTransport(hostname string, timeout time.Duration) (Transport, error) {
	// initialize ZMQ Context
	context, err := zmq.NewContext()
	if err != nil {
		return nil, err
	}

	// setup router and bind() to tcp address for clients to connect to
	router_sock, err := context.NewSocket(zmq.ROUTER)
	if err != nil {
		return nil, err
	}
	err = router_sock.Bind("tcp://" + hostname)
	if err != nil {
		return nil, err
	}

	// setup dealer for
	dealer_sock, err := context.NewSocket(zmq.DEALER)
	if err != nil {
		return nil, err
	}
	err = dealer_sock.Bind("inproc://dendrite-zmqdealer")
	if err != nil {
		return nil, err
	}
	poller := zmq.NewPoller()
	poller.Add(router_sock, zmq.POLLIN)
	poller.Add(dealer_sock, zmq.POLLIN)

	transport := &ZMQTransport{
		lock:              new(sync.Mutex),
		minHandlers:       10,
		maxHandlers:       1024,
		incrHandlers:      10,
		activeRequests:    0,
		workerIdleTimeout: 10 * time.Second,
		control_c:         make(chan *workerComm),
		dealer_sock:       dealer_sock,
		router_sock:       router_sock,
		zmq_context:       context,
	}
	// proxy messages between router and dealer
	go func() {
		for {
			sockets, _ := poller.Poll(-1)
			for _, socket := range sockets {
				switch s := socket.Socket; s {
				case router_sock:
					msg, err := s.RecvBytes(0)
					if err != nil {
						log.Println("ERR: TransportListener read bytes on router failed", err)
						continue
					}
					// forward to dealer
					_, err = dealer_sock.SendBytes(msg, 0)
					if err != nil {
						log.Println("ERR: TransportListener forward to dealer failed", err)
						continue
					}
					transport.activeRequests += 1
				case dealer_sock:
					msg, err := s.RecvBytes(0)
					if err != nil {
						log.Println("ERR: TransportListener read bytes on dealer failed", err)
						continue
					}
					// forward up to router
					_, err = router_sock.SendBytes(msg, 0)
					if err != nil {
						log.Println("ERR: TransportListener forward to router failed", err)
						continue
					}
					transport.activeRequests -= 1
				}
			}
		}
	}()

	// Scheduler goroutine keeps track of running workers
	// It spawns new ones if needed, and cancels ones that are idling
	go func() {
		sched_ticker := time.NewTicker(60 * time.Second)
		workers := make(map[*workerComm]bool)
		// fire up initial set of workers
		//for i := 0; i < transport.minHandlers; i++ {
		for i := 0; i < 15; i++ {
			//go handleZMQreq(transport)
		}
		for {
			select {
			case comm := <-transport.control_c:
				// worker sent something...
				msg := <-comm.worker_out
				switch {
				case msg == workerRegisterReq:
					if len(workers) == transport.maxHandlers {
						comm.worker_in <- workerRegisterDenied
						log.Println("ERR: TransportListener - max number of workers reached")
						continue
					}
					if _, ok := workers[comm]; ok {
						// worker already registered
						continue
					}
					comm.worker_in <- workerRegisterAllowed
					workers[comm] = true
					log.Println("INFO: TransportListener - registered new worker, total:", len(workers))
				case msg == workerShutdownReq:
					//log.Println("Got shutdown req")
					if len(workers) > transport.minHandlers {
						comm.worker_in <- workerShutdownAllowed
						for _ = range comm.worker_out {
							// wait until worker closes the channel
						}
						delete(workers, comm)
					} else {
						comm.worker_in <- workerShutdownDenied
					}
				}
			case <-sched_ticker.C:
				// check if requests are piling up and start more workers if that's the case
				if transport.activeRequests > 3*len(workers) {
					for i := 0; i < transport.incrHandlers; i++ {
						go handleZMQreq(transport)
					}
				}
			}
		}
	}()
	return transport, nil
}

type workerComm struct {
	worker_in  chan controlType // worker's input channel for two way communication with scheduler
	worker_out chan controlType // worker's output channel for two way communication with scheduler
	worker_ctl chan controlType // worker's control channel for communication with scheduler
}

func handleZMQreq(transport *ZMQTransport) {
	// setup REP socket
	rep_sock, err := transport.zmq_context.NewSocket(zmq.REP)
	if err != nil {
		log.Println("ERR: TransportListener worker failed to create REP socket", err)
		return
	}
	err = rep_sock.Connect("inproc://dendrite-zmqdealer")
	if err != nil {
		log.Println("ERR: TransportListener worker failed to connect to dealer", err)
		return
	}
{'fraction_non_alphanumeric': 0.09679878048780488, 'fraction_numerical': 0.005335365853658537, 'mean_word_length': 2.919716206123973, 'pattern_counts': {'":': 0, '<': 22, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7111616', 'n_tokens_mistral': 3934, 'n_tokens_neox': 3259, 'n_words': 971}