workerpool/example_test.go at main · alesr/workerpool · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
package workerpool_test

import (
	"context"
	"fmt"
	"os"
	"sync/atomic"

	"github.com/alesr/workerpool"
)

type bazooka struct {
	ammo      uint8
	targetID  string
	bodyCount *atomic.Int32
}

// Do simulate some bazooking
func (b *bazooka) Do(ctx context.Context) {
	b.ammo--
	fmt.Fprintln(os.Stderr, "bazooking: "+b.targetID)
	b.bodyCount.Add(1)
}

func Example_bounded() {
	pool := workerpool.New[*bazooka](context.TODO(), 3)

	var bodyCount atomic.Int32

	bazookas := []bazooka{
		{ammo: 69, targetID: "foo-id", bodyCount: &bodyCount},
		{ammo: 42, targetID: "bar-id", bodyCount: &bodyCount},
		{ammo: 11, targetID: "qux-id", bodyCount: &bodyCount},
	}

	for i := range bazookas {
		_ = pool.Submit(context.TODO(), &bazookas[i])
	}

	if err := pool.GracefulShutdown(); err != nil {
		fmt.Fprintf(os.Stderr, "shutdown error: %v\n", err)
	}

	fmt.Printf("Body count: %d\n", bodyCount.Load())
	// Output:
	// Body count: 3
}

func Example_unbounded() {
	pool := workerpool.New(context.TODO(), 3, workerpool.WithUnboundedQueue[*bazooka]())

	var bodyCount atomic.Int32

	bazookas := []bazooka{
		{ammo: 69, targetID: "foo-id", bodyCount: &bodyCount},
		{ammo: 42, targetID: "bar-id", bodyCount: &bodyCount},
		{ammo: 11, targetID: "qux-id", bodyCount: &bodyCount},
	}

	for i := range bazookas {
		_ = pool.Submit(context.TODO(), &bazookas[i])
	}

	if err := pool.GracefulShutdown(); err != nil {
		fmt.Fprintf(os.Stderr, "shutdown error: %v\n", err)
	}

	fmt.Printf("Body count: %d\n", bodyCount.Load())
	// Output:
	// Body count: 3
}