New `Compressor`: Support metadata other than `Name` for trained compression about grpc-go HOT 8 OPEN

// codec is a dictionary-aware compression implementation
//
// Given a category and namespace, it knows how to find the correct dictionary
// to compress with. For decompression, it reads the dictionary ID from the
// zstd payload.
//
// This could later be augmented to have decompression IDs category:namespace
// aware vs. assuming they're globally scoped.
type codec struct {
	Category   string
	Namespace string
}

func (c *codec) Marshal(v any) ([]byte, error) {
	return marshal(c.Category, c.Namespace, v)
}

func (c *codec) Unmarshal(data []byte, v any) error {
	return unmarshal(data, v)
}

func (c *codec) Name() string {
	return "zstd-wrapper"
}

// marshal and compress 'v', wrap in [zstd.Compressed], and return the result
// of marshalling that
//
// Skip wrapping if 'v' is already a [zstd.Compressed]
func marshal(category, namespace string, v any) ([]byte, error) {
	switch vv := v.(type) {
	case *zstd.Compressed:
		return proto.Marshal(vv)
	case proto.Message:
		data, err := proto.Marshal(vv)
		if err != nil {
			return nil, err
		}
		container := &zstd.Compressed{
			Category:   category,
			Namespace: namespace,
			// TODO: Imagine there was compression happening before this
			Data: data,
		}

		fmt.Println("using the codec for compression")
		return proto.Marshal(container)
	default:
		return nil, fmt.Errorf("unsure what to do for type: %T", v)
	}
}

// unmarshal data into a [zstd.Compressed], then unmarshal it's data field into 'v'
func unmarshal(data []byte, v any) error {
	switch vv := v.(type) {
	case *zstd.Compressed:
		return proto.Unmarshal(data, vv)
	case proto.Message:
		var container zstd.Compressed
		if err := proto.Unmarshal(data, &container); err != nil {
			return err
		}
		fmt.Println("using the codec for decompression")
		// TODO: Imagine there was decompression happening before this
		return proto.Unmarshal(container.Data, vv)
	default:
		return fmt.Errorf("unsure what to do for type: %T", v)
	}
}

package main

import (
	"context"
	"fmt"
	"log"
	"net"

	zstd "repro/gen"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/encoding"
	pb "google.golang.org/grpc/examples/helloworld/helloworld"
	"google.golang.org/protobuf/proto"
)

const addr = "localhost:50051"

type server struct {
	pb.UnimplementedGreeterServer
}

func (s *server) SayHello(ctx context.Context, in *pb.HelloRequest) (*pb.HelloReply, error) {
	return &pb.HelloReply{Message: "Hello " + in.GetName()}, nil
}

func main() {
	// Register default codec without category/namespace for decompression
	encoding.RegisterCodec(&codec{})

	lis, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}

	s := grpc.NewServer()
	pb.RegisterGreeterServer(s, &server{})

	go func() {
		if err := s.Serve(lis); err != nil {
			log.Fatalf("failed to serve: %v", err)
		}
	}()

	// Now connect from client
	conn, err := grpc.Dial(addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		log.Fatalf("did not connect: %v", err)
	}
	defer conn.Close()

	c := pb.NewGreeterClient(conn)
	r, err := c.SayHello(
		context.Background(),
		&pb.HelloRequest{Name: "coxley"},
		// TODO: This would ideally be cached and retrieved via an accessor somewhere
		grpc.ForceCodec(&codec{Category: "test", Namespace: "foo"}),
	)
	if err != nil {
		log.Fatalf("could not greet: %v", err)
	}
	fmt.Printf("Result: %v\n", r)
}

package grpc

func UseCompressor(name string, compressorOptions ...any) CallOption {}

func SetSendCompressor(ctx context.Context, name string, compressorOptions ...any) error {}

package encoding

type Compressor interface {
	Compress(/*see below*/, ...any) /*see below*/
	...
}