From de4a56b2f572a790d8c48b960a17b9a46b51eac6 Mon Sep 17 00:00:00 2001 From: Kyle Xiao Date: Fri, 19 Jul 2024 14:00:12 +0800 Subject: [PATCH] feat(thrift): port from kitex/pkg/utils/fastthrift --- go.mod | 5 +- go.sum | 10 ++++ protocol/thrift/fastcodec.go | 89 +++++++++++++++++++++++++++++++ protocol/thrift/fastcodec_test.go | 61 +++++++++++++++++++++ protocol/thrift/thrift.go | 15 ------ 5 files changed, 164 insertions(+), 16 deletions(-) create mode 100644 protocol/thrift/fastcodec.go create mode 100644 protocol/thrift/fastcodec_test.go diff --git a/go.mod b/go.mod index 8ad4bb4..db9c117 100644 --- a/go.mod +++ b/go.mod @@ -2,7 +2,10 @@ module github.com/cloudwego/gopkg go 1.17 -require github.com/stretchr/testify v1.9.0 +require ( + github.com/bytedance/gopkg v0.0.0-20240711085056-a03554c296f8 + github.com/stretchr/testify v1.9.0 +) require ( github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect diff --git a/go.sum b/go.sum index 340439d..c69d798 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +github.com/bytedance/gopkg v0.0.0-20240711085056-a03554c296f8 h1:rDwLxYTMoKHaw4cS0bQhaTZnkXp5e6ediCggGcRD/CA= +github.com/bytedance/gopkg v0.0.0-20240711085056-a03554c296f8/go.mod h1:FtQG3YbQG9L/91pbKSw787yBQPutC+457AvDW77fgUQ= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= @@ -13,11 +15,19 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+ github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +golang.org/x/net v0.0.0-20221014081412-f15817d10f9b/go.mod h1:YDH+HFinaLZZlnHAfSS6ZXJJ9M9t4Dl22yv3iI2vPwk= +golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220728004956-3c1f35247d10/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= +golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/protocol/thrift/fastcodec.go b/protocol/thrift/fastcodec.go new file mode 100644 index 0000000..69037cd --- /dev/null +++ b/protocol/thrift/fastcodec.go @@ -0,0 +1,89 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package thrift + +import ( + "errors" + + "github.com/bytedance/gopkg/lang/dirtmake" +) + +// NocopyWriteThreshold represents the threshold of using `NocopyWriter` for binary or string +// +// It's used by `WriteBinaryNocopy` and `WriteStringNocopy` of `BinaryProtocol` +// which are relied by kitex tool or thriftgo +var NocopyWriteThreshold = 4096 + +// BinaryWriter represents the method used in thrift encoding for nocopy writes +// It supports netpoll nocopy feature, see: https://github.com/cloudwego/netpoll/blob/develop/nocopy.go +type NocopyWriter interface { + WriteDirect(b []byte, remainCap int) error +} + +// FastCodec represents the interface of thrift fastcodec generated structs +type FastCodec interface { + BLength() int + FastWriteNocopy(buf []byte, bw NocopyWriter) int + FastRead(buf []byte) (int, error) +} + +// FastMarshal marshals the msg to buf. The msg should implement FastCodec. +func FastMarshal(msg FastCodec) []byte { + sz := msg.BLength() + buf := dirtmake.Bytes(sz, sz) + msg.FastWriteNocopy(buf, nil) + return buf +} + +// FastUnmarshal unmarshal the buf into msg. The msg should implement FastCodec. +func FastUnmarshal(buf []byte, msg FastCodec) error { + _, err := msg.FastRead(buf) + return err +} + +// MarshalFastMsg encodes the given msg to buf for generic thrift RPC. +func MarshalFastMsg(method string, msgType TMessageType, seq int32, msg FastCodec) ([]byte, error) { + if method == "" { + return nil, errors.New("method not set") + } + sz := Binary.MessageBeginLength(method, msgType, seq) + msg.BLength() + b := dirtmake.Bytes(sz, sz) + i := Binary.WriteMessageBegin(b, method, msgType, seq) + _ = msg.FastWriteNocopy(b[i:], nil) + return b, nil +} + +// UnmarshalFastMsg parses the given buf and stores the result to msg for generic thrift RPC. +// for EXCEPTION msgType, it will returns `err` with *ApplicationException type without storing the result to msg. +func UnmarshalFastMsg(b []byte, msg FastCodec) (method string, seq int32, err error) { + method, msgType, seq, i, err := Binary.ReadMessageBegin(b) + if err != nil { + return "", 0, err + } + b = b[i:] + + if msgType == EXCEPTION { + ex := NewApplicationException(UNKNOWN_APPLICATION_EXCEPTION, "") + _, err = ex.FastRead(b) + if err != nil { + return method, seq, err + } + return method, seq, ex + } + _, err = msg.FastRead(b) + return method, seq, err +} diff --git a/protocol/thrift/fastcodec_test.go b/protocol/thrift/fastcodec_test.go new file mode 100644 index 0000000..a6a8d68 --- /dev/null +++ b/protocol/thrift/fastcodec_test.go @@ -0,0 +1,61 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package thrift + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestFastMarshal(t *testing.T) { + req1, req2 := NewApplicationException(1, "hello"), NewApplicationException(0, "") + buf := FastMarshal(req1) + err := FastUnmarshal(buf, req2) + require.NoError(t, err) + require.Equal(t, req1.t, req2.t) + require.Equal(t, req1.m, req2.m) +} + +func TestMarshalFastMsg(t *testing.T) { + // CALL and REPLY + + req := NewApplicationException(1, "hello") + b, err := MarshalFastMsg("Echo", CALL, 1, req) + require.NoError(t, err) + + resp := NewApplicationException(0, "") + method, seq, err := UnmarshalFastMsg(b, resp) + require.NoError(t, err) + require.Equal(t, "Echo", method) + require.Equal(t, int32(1), seq) + require.Equal(t, req.t, resp.t) + require.Equal(t, req.m, resp.m) + + // EXCEPTION + + ex := NewApplicationException(WRONG_METHOD_NAME, "Ex!") + b, err = MarshalFastMsg("ExMethod", EXCEPTION, 2, ex) + require.NoError(t, err) + method, seq, err = UnmarshalFastMsg(b, nil) + require.NotNil(t, err) + require.Equal(t, "ExMethod", method) + require.Equal(t, int32(2), seq) + e, ok := err.(*ApplicationException) + require.True(t, ok) + require.True(t, e.TypeID() == ex.TypeID() && e.Error() == ex.Error()) +} diff --git a/protocol/thrift/thrift.go b/protocol/thrift/thrift.go index 2f9ed78..f1559a9 100644 --- a/protocol/thrift/thrift.go +++ b/protocol/thrift/thrift.go @@ -59,18 +59,3 @@ const ( // for Write/ReadMessage msgVersionMask = 0xffff0000 msgTypeMask = 0x0000ffff // for TMessageType ) - -var NocopyWriteThreshold = 4096 // use NocopyWriter when binary or string > the value - -// BinaryWriter represents the method used in thrift encoding for nocopy writes -// It supports netpoll nocopy feature, see: https://github.com/cloudwego/netpoll/blob/develop/nocopy.go -type NocopyWriter interface { - WriteDirect(b []byte, remainCap int) error -} - -// ThriftFastCodec represents the interface of thrift fastcodec generated structs -type ThriftFastCodec interface { - BLength() int - FastWriteNocopy(buf []byte, bw NocopyWriter) int - FastRead(buf []byte) (int, error) -}