github.com/matrixorigin/matrixone@v1.2.0/pkg/common/morpc/codec_header.go (about)

     1  // Copyright 2021 - 2022 Matrix Origin
     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 implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package morpc
    16  
    17  import (
    18  	"context"
    19  	"io"
    20  	"time"
    21  
    22  	"github.com/fagongzi/goetty/v2/buf"
    23  	"github.com/matrixorigin/matrixone/pkg/pb/timestamp"
    24  	"github.com/matrixorigin/matrixone/pkg/txn/clock"
    25  	"github.com/matrixorigin/matrixone/pkg/util/trace"
    26  )
    27  
    28  type deadlineContextCodec struct {
    29  }
    30  
    31  func (hc *deadlineContextCodec) Encode(msg *RPCMessage, out *buf.ByteBuf) (int, error) {
    32  	if msg.Ctx == nil {
    33  		return 0, nil
    34  	}
    35  
    36  	v, err := msg.GetTimeoutFromContext()
    37  	if err != nil {
    38  		return 0, err
    39  	}
    40  	out.WriteInt64(int64(v))
    41  	return 8, nil
    42  }
    43  
    44  func (hc *deadlineContextCodec) Decode(msg *RPCMessage, data []byte) (int, error) {
    45  	if len(data) < 8 {
    46  		return 0, io.ErrShortBuffer
    47  	}
    48  
    49  	if msg.Ctx == nil {
    50  		msg.Ctx = context.Background()
    51  	}
    52  
    53  	msg.Ctx, msg.Cancel = context.WithTimeout(msg.Ctx, time.Duration(buf.Byte2Int64(data)))
    54  	return 8, nil
    55  }
    56  
    57  type traceCodec struct {
    58  }
    59  
    60  func (hc *traceCodec) Encode(msg *RPCMessage, out *buf.ByteBuf) (int, error) {
    61  	if msg.Ctx == nil {
    62  		return 0, nil
    63  	}
    64  
    65  	span := trace.SpanFromContext(msg.Ctx)
    66  	c := span.SpanContext()
    67  	n := c.Size()
    68  	out.MustWriteByte(byte(n))
    69  	idx, _ := setWriterIndexAfterGow(out, n)
    70  	c.MarshalTo(out.RawSlice(idx, idx+n))
    71  	return 1 + n, nil
    72  }
    73  
    74  func (hc *traceCodec) Decode(msg *RPCMessage, data []byte) (int, error) {
    75  	if len(data) < 1 {
    76  		return 0, io.ErrShortBuffer
    77  	}
    78  
    79  	if len(data) < int(data[0]) {
    80  		return 0, io.ErrShortBuffer
    81  	}
    82  
    83  	c := &trace.SpanContext{}
    84  	if err := c.Unmarshal(data[1 : 1+data[0]]); err != nil {
    85  		return 0, err
    86  	}
    87  
    88  	if msg.Ctx == nil {
    89  		msg.Ctx = context.Background()
    90  	}
    91  	msg.Ctx = trace.ContextWithSpanContext(msg.Ctx, *c)
    92  	return int(1 + data[0]), nil
    93  }
    94  
    95  type hlcCodec struct {
    96  	clock clock.Clock
    97  }
    98  
    99  func (hc *hlcCodec) Encode(msg *RPCMessage, out *buf.ByteBuf) (int, error) {
   100  	now, _ := hc.clock.Now()
   101  	out.WriteInt64(now.PhysicalTime)
   102  	out.WriteUint32(now.LogicalTime)
   103  	return 12, nil
   104  }
   105  
   106  func (hc *hlcCodec) Decode(msg *RPCMessage, data []byte) (int, error) {
   107  	if len(data) < 12 {
   108  		return 0, io.ErrShortBuffer
   109  	}
   110  	ts := timestamp.Timestamp{}
   111  	ts.PhysicalTime = buf.Byte2Int64(data)
   112  	ts.LogicalTime = buf.Byte2Uint32(data[8:])
   113  	hc.clock.Update(ts)
   114  	return 12, nil
   115  }