worker-pool-with-thread-affinity.go

package main

import (
	"fmt"
	"runtime"
	"sync"
)

// Job represents the work unit
type Job struct {
	ID       int
	Request  []byte
	response chan []byte
}

func makeJob(id int) Job {
	return Job{
		ID:       id,
		response: make(chan []byte),
	}
}

func worker(id int, jobs <-chan Job) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	for j := range jobs {
		fmt.Printf("worker %d started job %d\n", id, j.ID)
		j.response <- []byte(fmt.Sprintf("ahoj_%d_%d", id, j.ID))
	}
}

// Manager handles jobs that need to be run with OS thread affinity
type Manager struct {
	jobs chan<- Job
}

// NewManager creates manager with an associated worker pool
func NewManager(poolSize int) *Manager {
	jobs := make(chan Job)
	for w := 1; w <= poolSize; w++ {
		go worker(w, jobs)
	}
	return &Manager{
		jobs: jobs,
	}
}

// WaitFor dispatches job and waits for its result
func (m *Manager) WaitFor(idx int) []byte {
	job := makeJob(idx)
	m.jobs <- job
	return <-job.response
}

func main() {
	m := NewManager(4)
	var wg sync.WaitGroup
	for i := 0; i < 16; i++ {
		wg.Add(1)
		go func(idx int) {
			fmt.Printf("%s\n", string(m.WaitFor(idx)))
			wg.Done()
		}(i)
	}
	wg.Wait()
}

TODO: Channels should be closed.