cubefs

1
// Copyright 2023 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 unboundedchan
16

17
import "sync/atomic"
18

19
type UnboundedChan struct {
20
	bufCount int64
21
	In       chan<- V // for user to write data into, BUT NOT TO READ FROM! In channel can be "close" by user
22
	Out      <-chan V // for user to read data from, BUT NOT TO WRITE INTO!
23
	buffer   *RingBuffer
24
}
25

26
func (uc *UnboundedChan) Len() int {
27
	return len(uc.In) + uc.BufLen() + len(uc.Out)
28
}
29

30
func (uc *UnboundedChan) BufLen() int {
31
	return int(atomic.LoadInt64(&uc.bufCount))
32
}
33

34
func NewUnboundedChan(initCapacity int) *UnboundedChan {
35
	return NewUnboundedChanSize(initCapacity, initCapacity, initCapacity)
36
}
37

38
func NewUnboundedChanSize(initInCapacity, initOutCapacity, initBufferCapacity int) *UnboundedChan {
39
	in := make(chan V, initInCapacity)
40
	out := make(chan V, initOutCapacity)
41
	uc := &UnboundedChan{
42
		In:     in,
43
		Out:    out,
44
		buffer: NewRingBuffer(uint32(initBufferCapacity)),
45
	}
46

47
	go run(in, out, uc)
48
	return uc
49
}
50

51
func (uc *UnboundedChan) feedBuffer(val V) {
52
	uc.buffer.Write(val)
53
	atomic.AddInt64(&uc.bufCount, 1)
54
}
55

56
func (uc *UnboundedChan) resetBuffer() {
57
	uc.buffer.Reset()
58
	atomic.StoreInt64(&uc.bufCount, 0)
59
}
60

61
func run(in, out chan V, uc *UnboundedChan) {
62
	defer close(out)
63
LOOP:
64
	for {
65
		// read data from in channel
66
		val, ok := <-in
67
		if !ok {
68
			break LOOP
69
		}
70

71
		// move data to buffer or out channel
72
		if atomic.LoadInt64(&uc.bufCount) > 0 {
73
			uc.feedBuffer(val)
74
		} else {
75
			select {
76
			case out <- val:
77
				continue
78
			default:
79
				// out channel is full
80
				uc.feedBuffer(val)
81
			}
82
		}
83

84
		// try feeding out channel with buffer data
85
		for !uc.buffer.IsEmpty() {
86
			select {
87
			// when out channel is full, data may still feed in channel
88
			case val, ok := <-in:
89
				if !ok {
90
					break LOOP
91
				}
92
				uc.feedBuffer(val)
93
			case out <- uc.buffer.Peek():
94
				uc.buffer.Pop()
95
				atomic.AddInt64(&uc.bufCount, -1)
96
				if uc.buffer.IsEmpty() {
97
					uc.resetBuffer()
98
				}
99
			}
100
		}
101
	}
102

103
	// no more in data, keep feeding out channel with buffer data until it's drained
104
	for !uc.buffer.IsEmpty() {
105
		out <- uc.buffer.Pop()
106
		atomic.AddInt64(&uc.bufCount, -1)
107
	}
108
	uc.resetBuffer()
109
}
110