Add multi worker pool

sync/example_test.go

@@ -89,3 +89,93 @@ func ExampleWorkerPool() {
 	fmt.Printf("Total is %v", total)
 	// Output: Total is 36
 }
+
+func ExampleMultiWorkerPool() {
+	var (
+		errorCh      = make(chan error, 1)
+		processWg    sync.WaitGroup
+		workers      = xsync.NewMultiWorkerPool(2, 3)
+		numProcesses = 3
+		numRequests  = 9
+		responses    = make([][]response, numProcesses)
+	)
+
+	for i := 0; i < numProcesses; i++ {
+		responses[i] = make([]response, numRequests)
+	}
+
+	// Initialise the multi worker pool
+	workers.Init()
+	processWg.Add(numProcesses)
+
+	for process := 0; process < numProcesses; process++ {
+		//Capture loop variable.
+		process := process
+
+		// Start a request per process
+		go func() {
+			// Get a worker pool
+			pool := workers.GetPool()
+
+			// Spin up a wait group per process
+			var wg sync.WaitGroup
+			wg.Add(numRequests)
+			for i := 0; i < numRequests; i++ {
+				// Capture loop variable.
+				i := i
+
+				// Execute request on worker pool.
+				pool.Go(func() {
+					defer wg.Done()
+
+					var err error
+
+					// Perform some work which may fail.
+					resp := response{a: i * (process + 1)}
+
+					if err != nil {
+						// Return the first error that is encountered.
+						select {
+						case errorCh <- err:
+						default:
+						}
+
+						return
+					}
+
+					// Can concurrently modify responses since each iteration updates a
+					// different index.
+					responses[process][i] = resp
+				})
+			}
+			// Wait for all requests to finish.
+			wg.Wait()
+
+			// Return pool to worker pool
+			workers.PutPool(pool)
+
+			// Wait for all process requests to finish.
+			processWg.Done()
+		}()
+	}
+
+	// Wait for all process requests to finish.
+	processWg.Wait()
+
+	close(errorCh)
+	if err := <-errorCh; err != nil {
+		log.Fatal(err)
+	}
+
+	for i, r := range responses {
+		total := 0
+		for _, response := range r {
+			total += response.a
+		}
+		fmt.Printf("Total for process %v is %v\n", i, total)
+	}
+
+	// Output: Total for process 0 is 36
+	// Total for process 1 is 72
+	// Total for process 2 is 108
+}

sync/multi_worker_pool.go

@@ -0,0 +1,65 @@
+// Copyright (c) 2018 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+// Package sync implements synchronization facililites such as worker pools.
+package sync
+
+// MultiWorkerPool provides a pool for WorkerPools. This can be useful if multiple
+// independent processes need to share a set amount of pools, but with bounding conditions
+// as to how many workers a single process can consume
+type MultiWorkerPool interface {
+	// Init initializes the pool.
+	Init()
+
+	// GetPool provides a WorkerPool
+	GetPool() WorkerPool
+
+	// PutPool returns a WorkerPool to the pool
+	PutPool(pool WorkerPool)
+}
+
+type multiWorkerPool struct {
+	poolCh   chan WorkerPool
+	poolSize int
+}
+
+// NewMultiWorkerPool creates a new multi worker pool.
+func NewMultiWorkerPool(size int, poolSize int) MultiWorkerPool {
+	return &multiWorkerPool{
+		poolCh:   make(chan WorkerPool, size),
+		poolSize: poolSize,
+	}
+}
+
+func (p *multiWorkerPool) Init() {
+	for i := 0; i < cap(p.poolCh); i++ {
+		innerPool := NewWorkerPool(p.poolSize)
+		innerPool.Init()
+		p.poolCh <- innerPool
+	}
+}
+
+func (p *multiWorkerPool) GetPool() WorkerPool {
+	return <-p.poolCh
+}
+
+func (p *multiWorkerPool) PutPool(pool WorkerPool) {
+	p.poolCh <- pool
+}

sync/multi_worker_pool_test.go

@@ -0,0 +1,60 @@
+// Copyright (c) 2017 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package sync
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+const multiTestWorkerPoolSize = 2
+
+func TestGetRoutine(t *testing.T) {
+	var count uint32
+
+	multi := NewMultiWorkerPool(multiTestWorkerPoolSize, testWorkerPoolSize)
+	multi.Init()
+
+	var multiWg sync.WaitGroup
+	for pool := 0; pool < multiTestWorkerPoolSize*2; pool++ {
+		multiWg.Add(1)
+		go func() {
+			p := multi.GetPool()
+			var wg sync.WaitGroup
+			for i := 0; i < testWorkerPoolSize*2; i++ {
+				wg.Add(1)
+				p.Go(func() {
+					atomic.AddUint32(&count, 1)
+					wg.Done()
+				})
+			}
+			wg.Wait()
+			multi.PutPool(p)
+			multiWg.Done()
+		}()
+	}
+	multiWg.Wait()
+
+	require.Equal(t, uint32(testWorkerPoolSize*2*multiTestWorkerPoolSize*2), count)
+}