github.com/cloudwego/kitex@v0.9.0/pkg/remote/payload_codec.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package remote
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  
    23  	"github.com/cloudwego/kitex/pkg/serviceinfo"
    24  )
    25  
    26  var payloadCodecs = make(map[serviceinfo.PayloadCodec]PayloadCodec)
    27  
    28  // PayloadCodec is used to marshal and unmarshal payload.
    29  type PayloadCodec interface {
    30  	Marshal(ctx context.Context, message Message, out ByteBuffer) error
    31  	Unmarshal(ctx context.Context, message Message, in ByteBuffer) error
    32  	Name() string
    33  }
    34  
    35  // GetPayloadCodec gets desired payload codec from message.
    36  func GetPayloadCodec(message Message) (PayloadCodec, error) {
    37  	if message.PayloadCodec() != nil {
    38  		return message.PayloadCodec(), nil
    39  	}
    40  	ct := message.ProtocolInfo().CodecType
    41  	pc := payloadCodecs[ct]
    42  	if pc == nil {
    43  		return nil, fmt.Errorf("payload codec not found with codecType=%v", ct)
    44  	}
    45  	return pc, nil
    46  }
    47  
    48  // PutPayloadCode puts the desired payload codec to message.
    49  func PutPayloadCode(name serviceinfo.PayloadCodec, v PayloadCodec) {
    50  	payloadCodecs[name] = v
    51  }