WIP: Use proto if we have proto types

Author: justinsb

pkg/cache/internal/informers_map.go

@@ -266,7 +266,7 @@ func createStructuredListWatch(gvk schema.GroupVersionKind, ip *specificInformer
 		return nil, err
 	}
 
-	client, err := apiutil.RESTClientForGVK(gvk, false, ip.config, ip.codecs)
+	client, err := apiutil.RESTClientForGVK(ip.Scheme, gvk, false, ip.config, ip.codecs)
 	if err != nil {
 		return nil, err
 	}

pkg/client/apiutil/apimachinery.go

@@ -21,41 +21,18 @@ package apiutil
 
 import (
 	"fmt"
-	"sync"
+	"reflect"
 
 	"k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/runtime/serializer"
 	"k8s.io/client-go/discovery"
-	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/restmapper"
 )
 
-var (
-	protobufScheme     = runtime.NewScheme()
-	protobufSchemeLock sync.RWMutex
-)
-
-func init() {
-	// Currently only enabled for built-in resources which are guaranteed to implement Protocol Buffers.
-	// For custom resources, CRDs can not support Protocol Buffers but Aggregated API can.
-	// See doc: https://kubernetes.io/docs/concepts/extend-kubernetes/api-extension/custom-resources/#advanced-features-and-flexibility
-	if err := clientgoscheme.AddToScheme(protobufScheme); err != nil {
-		panic(err)
-	}
-}
-
-// AddToProtobufScheme add the given SchemeBuilder into protobufScheme, which should
-// be additional types that do support protobuf.
-func AddToProtobufScheme(addToScheme func(*runtime.Scheme) error) error {
-	protobufSchemeLock.Lock()
-	defer protobufSchemeLock.Unlock()
-	return addToScheme(protobufScheme)
-}
-
 // NewDiscoveryRESTMapper constructs a new RESTMapper based on discovery
 // information fetched by a new client with the given config.
 func NewDiscoveryRESTMapper(c *rest.Config) (meta.RESTMapper, error) {
@@ -117,8 +94,8 @@ func GVKForObject(obj runtime.Object, scheme *runtime.Scheme) (schema.GroupVersi
 // RESTClientForGVK constructs a new rest.Interface capable of accessing the resource associated
 // with the given GroupVersionKind. The REST client will be configured to use the negotiated serializer from
 // baseConfig, if set, otherwise a default serializer will be set.
-func RESTClientForGVK(gvk schema.GroupVersionKind, isUnstructured bool, baseConfig *rest.Config, codecs serializer.CodecFactory) (rest.Interface, error) {
-	return rest.RESTClientFor(createRestConfig(gvk, isUnstructured, baseConfig, codecs))
+func RESTClientForGVK(scheme *runtime.Scheme, gvk schema.GroupVersionKind, isUnstructured bool, baseConfig *rest.Config, codecs serializer.CodecFactory) (rest.Interface, error) {
+	return rest.RESTClientFor(createRestConfig(scheme, gvk, isUnstructured, baseConfig, codecs))
 }
 
 // serializerWithDecodedGVK is a CodecFactory that overrides the DecoderToVersion of a WithoutConversionCodecFactory
@@ -135,7 +112,7 @@ func (f serializerWithDecodedGVK) DecoderToVersion(serializer runtime.Decoder, _
 }
 
 // createRestConfig copies the base config and updates needed fields for a new rest config.
-func createRestConfig(gvk schema.GroupVersionKind, isUnstructured bool, baseConfig *rest.Config, codecs serializer.CodecFactory) *rest.Config {
+func createRestConfig(scheme *runtime.Scheme, gvk schema.GroupVersionKind, isUnstructured bool, baseConfig *rest.Config, codecs serializer.CodecFactory) *rest.Config {
 	gv := gvk.GroupVersion()
 
 	cfg := rest.CopyConfig(baseConfig)
@@ -150,11 +127,9 @@ func createRestConfig(gvk schema.GroupVersionKind, isUnstructured bool, baseConf
 	}
 	// TODO(FillZpp): In the long run, we want to check discovery or something to make sure that this is actually true.
 	if cfg.ContentType == "" && !isUnstructured {
-		protobufSchemeLock.RLock()
-		if protobufScheme.Recognizes(gvk) {
+		if haveProtoDefinition(scheme, gvk) {
 			cfg.ContentType = runtime.ContentTypeProtobuf
 		}
-		protobufSchemeLock.RUnlock()
 	}
 
 	if cfg.NegotiatedSerializer == nil {
@@ -169,3 +144,23 @@ func createRestConfig(gvk schema.GroupVersionKind, isUnstructured bool, baseConf
 
 	return cfg
 }
+
+// protoMessage is implemented by protobuf messages (of all libraries)
+type protoMessage interface {
+	ProtoMessage()
+}
+
+var protoMessageType = reflect.TypeOf(new(protoMessage)).Elem()
+
+// haveProtoDefinition returns true if the go type for the specified gvk support protobuf encoding.
+func haveProtoDefinition(scheme *runtime.Scheme, gvk schema.GroupVersionKind) bool {
+	gvkType, found := scheme.AllKnownTypes()[gvk]
+	if !found {
+		return false
+	}
+	// We normally get the raw struct e.g. v1.Pod, not &v1.Pod
+	if gvkType.Kind() == reflect.Struct {
+		gvkType = reflect.PtrTo(gvkType)
+	}
+	return gvkType.Implements(protoMessageType)
+}

pkg/client/apiutil/apimachinery_test.go

@@ -0,0 +1,117 @@
+/*
+Copyright 2021 The Kubernetes 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 apiutil contains utilities for working with raw Kubernetes
+// API machinery, such as creating RESTMappers and raw REST clients,
+// and extracting the GVK of an object.
+package apiutil
+
+import (
+	"fmt"
+	"testing"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/schema"
+	kubernetesscheme "k8s.io/client-go/kubernetes/scheme"
+	"sigs.k8s.io/controller-runtime/pkg/scheme"
+)
+
+var exampleSchemeGroupVersion = schema.GroupVersion{Group: "example.com", Version: "v1"}
+
+type ExampleCRD struct {
+	metav1.TypeMeta   `json:",inline"`
+	metav1.ObjectMeta `json:"metadata,omitempty"`
+}
+
+func (e *ExampleCRD) DeepCopyObject() runtime.Object {
+	panic("not implemented")
+}
+
+func TestHaveProtoDefinition(t *testing.T) {
+	exampleCRDScheme := runtime.NewScheme()
+
+	builder := &scheme.Builder{GroupVersion: exampleSchemeGroupVersion}
+	builder.Register(&ExampleCRD{})
+	if err := builder.AddToScheme(exampleCRDScheme); err != nil {
+		t.Fatalf("AddToScheme failed: %v", err)
+	}
+
+	schemes := map[string]*runtime.Scheme{
+		"kubernetes":  kubernetesscheme.Scheme,
+		"empty":       runtime.NewScheme(),
+		"example.com": exampleCRDScheme,
+	}
+	grid := []struct {
+		scheme                  string
+		gvk                     schema.GroupVersionKind
+		wantType                string
+		wantHaveProtoDefinition bool
+	}{
+		{
+			scheme:                  "kubernetes",
+			gvk:                     schema.GroupVersionKind{Group: "", Version: "v1", Kind: "Pod"},
+			wantType:                "v1.Pod",
+			wantHaveProtoDefinition: true,
+		},
+		{
+			scheme:                  "kubernetes",
+			gvk:                     schema.GroupVersionKind{Group: "example.com", Version: "v1", Kind: "Pod"},
+			wantType:                "",
+			wantHaveProtoDefinition: false,
+		},
+		{
+			scheme:                  "empty",
+			gvk:                     schema.GroupVersionKind{Group: "", Version: "v1", Kind: "Pod"},
+			wantType:                "",
+			wantHaveProtoDefinition: false,
+		},
+		{
+			scheme:                  "example.com",
+			gvk:                     schema.GroupVersionKind{Group: "", Version: "v1", Kind: "Pod"},
+			wantType:                "",
+			wantHaveProtoDefinition: false,
+		},
+		{
+			scheme:                  "example.com",
+			gvk:                     schema.GroupVersionKind{Group: "example.com", Version: "v1", Kind: "ExampleCRD"},
+			wantType:                "apiutil.ExampleCRD",
+			wantHaveProtoDefinition: false,
+		},
+	}
+
+	for _, g := range grid {
+		t.Run(fmt.Sprintf("%#v", g), func(t *testing.T) {
+			scheme := schemes[g.scheme]
+			if scheme == nil {
+				t.Errorf("scheme %q not found", g.scheme)
+			}
+
+			gotType := ""
+			if t := scheme.AllKnownTypes()[g.gvk]; t != nil {
+				gotType = t.String()
+			}
+			if gotType != g.wantType {
+				t.Errorf("unexpected type got %v, want %v", gotType, g.wantType)
+
+			}
+			gotHaveProtoDefinition := haveProtoDefinition(scheme, g.gvk)
+			if gotHaveProtoDefinition != g.wantHaveProtoDefinition {
+				t.Errorf("haveProtoDefintion(%#v, %#v) got %v, want %v", g.scheme, g.gvk, gotHaveProtoDefinition, g.wantHaveProtoDefinition)
+			}
+		})
+	}
+}

pkg/client/client_cache.go

@@ -59,7 +59,7 @@ func (c *clientCache) newResource(gvk schema.GroupVersionKind, isList, isUnstruc
 		gvk.Kind = gvk.Kind[:len(gvk.Kind)-4]
 	}
 
-	client, err := apiutil.RESTClientForGVK(gvk, isUnstructured, c.config, c.codecs)
+	client, err := apiutil.RESTClientForGVK(c.scheme, gvk, isUnstructured, c.config, c.codecs)
 	if err != nil {
 		return nil, err
 	}