google.golang.org/grpc@v1.72.2/encoding/proto/proto.go (about)

     1  /*
     2   *
     3   * Copyright 2024 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  // Package proto defines the protobuf codec. Importing this package will
    20  // register the codec.
    21  package proto
    22  
    23  import (
    24  	"fmt"
    25  
    26  	"google.golang.org/grpc/encoding"
    27  	"google.golang.org/grpc/mem"
    28  	"google.golang.org/protobuf/proto"
    29  	"google.golang.org/protobuf/protoadapt"
    30  )
    31  
    32  // Name is the name registered for the proto compressor.
    33  const Name = "proto"
    34  
    35  func init() {
    36  	encoding.RegisterCodecV2(&codecV2{})
    37  }
    38  
    39  // codec is a CodecV2 implementation with protobuf. It is the default codec for
    40  // gRPC.
    41  type codecV2 struct{}
    42  
    43  func (c *codecV2) Marshal(v any) (data mem.BufferSlice, err error) {
    44  	vv := messageV2Of(v)
    45  	if vv == nil {
    46  		return nil, fmt.Errorf("proto: failed to marshal, message is %T, want proto.Message", v)
    47  	}
    48  
    49  	size := proto.Size(vv)
    50  	if mem.IsBelowBufferPoolingThreshold(size) {
    51  		buf, err := proto.Marshal(vv)
    52  		if err != nil {
    53  			return nil, err
    54  		}
    55  		data = append(data, mem.SliceBuffer(buf))
    56  	} else {
    57  		pool := mem.DefaultBufferPool()
    58  		buf := pool.Get(size)
    59  		if _, err := (proto.MarshalOptions{}).MarshalAppend((*buf)[:0], vv); err != nil {
    60  			pool.Put(buf)
    61  			return nil, err
    62  		}
    63  		data = append(data, mem.NewBuffer(buf, pool))
    64  	}
    65  
    66  	return data, nil
    67  }
    68  
    69  func (c *codecV2) Unmarshal(data mem.BufferSlice, v any) (err error) {
    70  	vv := messageV2Of(v)
    71  	if vv == nil {
    72  		return fmt.Errorf("failed to unmarshal, message is %T, want proto.Message", v)
    73  	}
    74  
    75  	buf := data.MaterializeToBuffer(mem.DefaultBufferPool())
    76  	defer buf.Free()
    77  	// TODO: Upgrade proto.Unmarshal to support mem.BufferSlice. Right now, it's not
    78  	//  really possible without a major overhaul of the proto package, but the
    79  	//  vtprotobuf library may be able to support this.
    80  	return proto.Unmarshal(buf.ReadOnlyData(), vv)
    81  }
    82  
    83  func messageV2Of(v any) proto.Message {
    84  	switch v := v.(type) {
    85  	case protoadapt.MessageV1:
    86  		return protoadapt.MessageV2Of(v)
    87  	case protoadapt.MessageV2:
    88  		return v
    89  	}
    90  
    91  	return nil
    92  }
    93  
    94  func (c *codecV2) Name() string {
    95  	return Name
    96  }