Skip to content

Commit 099cb5b

Browse files
authored
Fixed buffered race condition (#136)
* Fixed buffered race condition * Fix go 1.18 compatibility * Bump go.uber.org/atomic to v1.10.0 * Use unsafe atomic pointer
1 parent f3f8376 commit 099cb5b

File tree

1 file changed

+44
-23
lines changed

1 file changed

+44
-23
lines changed

core/framing/buffered.go

+44-23
Original file line numberDiff line numberDiff line change
@@ -4,91 +4,112 @@ import (
44
"encoding/binary"
55
"io"
66
"sync/atomic"
7+
"unsafe"
78

89
"github.com/rsocket/rsocket-go/core"
910
"github.com/rsocket/rsocket-go/internal/common"
1011
"github.com/rsocket/rsocket-go/internal/u24"
12+
uberatomic "go.uber.org/atomic"
1113
)
1214

1315
// bufferedFrame is basic frame implementation.
1416
type bufferedFrame struct {
15-
inner *common.ByteBuff
16-
refs int32
17+
innerPtr unsafe.Pointer
18+
refs uberatomic.Int32
1719
}
1820

1921
func newBufferedFrame(inner *common.ByteBuff) *bufferedFrame {
20-
return &bufferedFrame{
21-
inner: inner,
22-
refs: 1,
23-
}
22+
frame := &bufferedFrame{}
23+
atomic.StorePointer(&frame.innerPtr, unsafe.Pointer(inner))
24+
frame.refs.Store(1)
25+
return frame
2426
}
2527

2628
func (f *bufferedFrame) IncRef() int32 {
27-
return atomic.AddInt32(&f.refs, 1)
29+
return f.refs.Add(1)
2830
}
2931

3032
func (f *bufferedFrame) RefCnt() int32 {
31-
return atomic.LoadInt32(&f.refs)
33+
return f.refs.Load()
3234
}
3335

3436
func (f *bufferedFrame) Header() core.FrameHeader {
35-
if f.inner == nil {
37+
inner := (*common.ByteBuff)(atomic.LoadPointer(&f.innerPtr))
38+
if inner == nil {
3639
panic("frame has been released!")
3740
}
38-
b := f.inner.Bytes()
41+
b := inner.Bytes()
3942
_ = b[core.FrameHeaderLen-1]
4043
var h core.FrameHeader
4144
copy(h[:], b)
4245
return h
4346
}
4447

4548
func (f *bufferedFrame) HasFlag(flag core.FrameFlag) bool {
46-
if f.inner == nil {
49+
inner := (*common.ByteBuff)(atomic.LoadPointer(&f.innerPtr))
50+
if inner == nil {
4751
panic("frame has been released!")
4852
}
49-
n := binary.BigEndian.Uint16(f.inner.Bytes()[4:6])
53+
n := binary.BigEndian.Uint16(inner.Bytes()[4:6])
5054
return core.FrameFlag(n&0x03FF)&flag == flag
5155
}
5256

5357
func (f *bufferedFrame) StreamID() uint32 {
54-
if f.inner == nil {
58+
inner := (*common.ByteBuff)(atomic.LoadPointer(&f.innerPtr))
59+
if inner == nil {
5560
panic("frame has been released!")
5661
}
57-
return binary.BigEndian.Uint32(f.inner.Bytes()[:4])
62+
return binary.BigEndian.Uint32(inner.Bytes()[:4])
5863
}
5964

6065
// Release releases resource.
6166
func (f *bufferedFrame) Release() {
62-
if f != nil && f.inner != nil && atomic.AddInt32(&f.refs, -1) == 0 {
63-
common.ReturnByteBuff(f.inner)
64-
f.inner = nil
67+
if f == nil {
68+
return
69+
}
70+
refs := f.refs.Add(-1)
71+
if refs > 0 {
72+
return
73+
}
74+
inner := (*common.ByteBuff)(atomic.LoadPointer(&f.innerPtr))
75+
if inner != nil {
76+
swapped := atomic.CompareAndSwapPointer(&f.innerPtr, unsafe.Pointer(inner), unsafe.Pointer(nil))
77+
if swapped {
78+
common.ReturnByteBuff(inner)
79+
}
6580
}
6681
}
6782

6883
// Body returns frame body.
6984
func (f *bufferedFrame) Body() []byte {
70-
if f.inner == nil {
85+
inner := (*common.ByteBuff)(atomic.LoadPointer(&f.innerPtr))
86+
if inner == nil {
7187
return nil
7288
}
73-
b := f.inner.Bytes()
89+
b := inner.Bytes()
7490
_ = b[core.FrameHeaderLen-1]
7591
return b[core.FrameHeaderLen:]
7692
}
7793

7894
// Len returns length of frame.
7995
func (f *bufferedFrame) Len() int {
80-
if f.inner == nil {
96+
inner := (*common.ByteBuff)(atomic.LoadPointer(&f.innerPtr))
97+
if inner == nil {
8198
return 0
8299
}
83-
return f.inner.Len()
100+
return inner.Len()
84101
}
85102

86103
// WriteTo write frame to writer.
87104
func (f *bufferedFrame) WriteTo(w io.Writer) (n int64, err error) {
88-
if f == nil || f.inner == nil {
105+
if f == nil {
106+
return
107+
}
108+
inner := (*common.ByteBuff)(atomic.LoadPointer(&f.innerPtr))
109+
if inner == nil {
89110
return
90111
}
91-
n, err = f.inner.WriteTo(w)
112+
n, err = inner.WriteTo(w)
92113
return
93114
}
94115

0 commit comments

Comments
 (0)