/
module_compression.go
139 lines (115 loc) · 3.42 KB
/
module_compression.go
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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
package codec
import (
"bytes"
"errors"
"fmt"
"git.golaxy.org/core"
"git.golaxy.org/framework/net/gtp"
"git.golaxy.org/framework/net/gtp/method"
"git.golaxy.org/framework/util/binaryutil"
"io"
"math"
)
// ICompressionModule 压缩模块接口
type ICompressionModule interface {
// Compress 压缩数据
Compress(src []byte) (dst binaryutil.RecycleBytes, compressed bool, err error)
// Uncompress 解压缩数据
Uncompress(src []byte) (dst binaryutil.RecycleBytes, err error)
}
// NewCompressionModule 创建压缩模块
func NewCompressionModule(cs method.CompressionStream) ICompressionModule {
if cs == nil {
panic(fmt.Errorf("%w: cs is nil", core.ErrArgs))
}
return &CompressionModule{
CompressionStream: cs,
}
}
// CompressionModule 压缩模块
type CompressionModule struct {
CompressionStream method.CompressionStream // 压缩流
}
// Compress 压缩数据
func (m *CompressionModule) Compress(src []byte) (dst binaryutil.RecycleBytes, compressed bool, err error) {
if len(src) <= 0 {
return binaryutil.MakeNonRecycleBytes(src), false, nil
}
if m.CompressionStream == nil {
return binaryutil.MakeNonRecycleBytes(nil), false, errors.New("setting CompressionStream is nil")
}
compressedBuf := binaryutil.MakeRecycleBytes(binaryutil.BytesPool.Get(len(src)))
defer compressedBuf.Release()
n, err := func() (n int, err error) {
bw := binaryutil.NewBytesWriter(compressedBuf.Data())
w, err := m.CompressionStream.WrapWriter(bw)
if err != nil {
return 0, err
}
defer func() {
if err == nil {
if err = w.Close(); err == nil {
n = bw.N
}
} else {
w.Close()
}
}()
_, err = w.Write(src)
return
}()
if err != nil {
if errors.Is(err, binaryutil.ErrLimitReached) {
return binaryutil.MakeNonRecycleBytes(src), false, nil
}
return binaryutil.MakeNonRecycleBytes(nil), false, err
}
msgCompressed := gtp.MsgCompressed{
Data: compressedBuf.Data()[:n],
OriginalSize: int64(len(src)),
}
if msgCompressed.Size() >= len(src) {
return binaryutil.MakeNonRecycleBytes(src), false, nil
}
buf := binaryutil.MakeRecycleBytes(binaryutil.BytesPool.Get(msgCompressed.Size()))
defer func() {
if err != nil {
buf.Release()
}
}()
if _, err = msgCompressed.Read(buf.Data()); err != nil {
return binaryutil.MakeNonRecycleBytes(nil), false, err
}
return buf, true, nil
}
// Uncompress 解压缩数据
func (m *CompressionModule) Uncompress(src []byte) (dst binaryutil.RecycleBytes, err error) {
if len(src) <= 0 {
return binaryutil.MakeNonRecycleBytes(nil), fmt.Errorf("%w: src too small", core.ErrArgs)
}
if m.CompressionStream == nil {
return binaryutil.MakeNonRecycleBytes(nil), errors.New("setting CompressionStream is nil")
}
msgCompressed := gtp.MsgCompressed{}
_, err = msgCompressed.Write(src)
if err != nil {
return binaryutil.MakeNonRecycleBytes(nil), err
}
if msgCompressed.OriginalSize >= math.MaxInt32 {
return binaryutil.MakeNonRecycleBytes(nil), errors.New("original size too large")
}
rawBuf := binaryutil.MakeRecycleBytes(binaryutil.BytesPool.Get(int(msgCompressed.OriginalSize)))
defer func() {
if err != nil {
rawBuf.Release()
}
}()
r, err := m.CompressionStream.WrapReader(bytes.NewReader(msgCompressed.Data))
if err != nil {
return binaryutil.MakeNonRecycleBytes(nil), err
}
if _, err = r.Read(rawBuf.Data()); err != nil && !errors.Is(err, io.EOF) {
return binaryutil.MakeNonRecycleBytes(nil), err
}
return rawBuf, nil
}