cubefs

1
// Copyright 2022 The CubeFS Authors.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
12
// implied. See the License for the specific language governing
13
// permissions and limitations under the License.
14

15
// Package taskpool provides limited pool running task
16
package common
17

18
// TaskPool limited pool
19
type TaskPool struct {
20
	pool chan func()
21
}
22

23
// New returns task pool with workerCount and poolSize
24
func New(workerCount, poolSize int) TaskPool {
25
	pool := make(chan func(), poolSize)
26
	for i := 0; i < workerCount; i++ {
27
		go func() {
28
			for {
29
				task, ok := <-pool
30
				if !ok {
31
					break
32
				}
33
				task()
34
			}
35
		}()
36
	}
37
	return TaskPool{pool: pool}
38
}
39

40
// Run add task to pool, block if pool is full
41
func (tp TaskPool) Run(task func()) {
42
	tp.pool <- task
43
}
44

45
// TryRun try to add task to pool, return immediately
46
func (tp TaskPool) TryRun(task func()) bool {
47
	select {
48
	case tp.pool <- task:
49
		return true
50
	default:
51
		return false
52
	}
53
}
54

55
// Close the pool, the function is concurrent unsafe
56
func (tp TaskPool) Close() {
57
	close(tp.pool)
58
}
59