Add lazy marshaling to handler.go.

Author: DangerOnTheRanger

pkg/handler/handler.go

@@ -33,6 +33,7 @@ import (
 	jsoniter "github.com/json-iterator/go"
 	"github.com/munnerz/goautoneg"
 	"gopkg.in/yaml.v2"
+	klog "k8s.io/klog/v2"
 	"k8s.io/kube-openapi/pkg/builder"
 	"k8s.io/kube-openapi/pkg/common"
 	"k8s.io/kube-openapi/pkg/validation/spec"
@@ -55,13 +56,26 @@ type OpenAPIService struct {
 
 	lastModified time.Time
 
-	specBytes []byte
-	specPb    []byte
-	specPbGz  []byte
+	jsonCache  cache
+	protoCache cache
+}
+
+type cache struct {
+	BuildCache func() ([]byte, error)
+	once       sync.Once
+	bytes      []byte
+	etag       string
+	err        error
+}
 
-	specBytesETag string
-	specPbETag    string
-	specPbGzETag  string
+func (c *cache) Get() ([]byte, string, error) {
+	c.once.Do(func() {
+		c.bytes, c.err = c.BuildCache()
+		if c.err != nil {
+			c.etag = computeETag(c.bytes)
+		}
+	})
+	return c.bytes, c.etag, c.err
 }
 
 func init() {
@@ -83,50 +97,44 @@ func NewOpenAPIService(spec *spec.Swagger) (*OpenAPIService, error) {
 	return o, nil
 }
 
-func (o *OpenAPIService) getSwaggerBytes() ([]byte, string, time.Time) {
-	o.rwMutex.RLock()
-	defer o.rwMutex.RUnlock()
-	return o.specBytes, o.specBytesETag, o.lastModified
-}
-
-func (o *OpenAPIService) getSwaggerPbBytes() ([]byte, string, time.Time) {
+func (o *OpenAPIService) getSwaggerBytes() ([]byte, string, time.Time, error) {
 	o.rwMutex.RLock()
 	defer o.rwMutex.RUnlock()
-	return o.specPb, o.specPbETag, o.lastModified
+	specBytes, etag, err := o.jsonCache.Get()
+	if err != nil {
+		return nil, "", time.Time{}, err
+	}
+	return specBytes, etag, o.lastModified, nil
 }
 
-func (o *OpenAPIService) getSwaggerPbGzBytes() ([]byte, string, time.Time) {
+func (o *OpenAPIService) getSwaggerPbBytes() ([]byte, string, time.Time, error) {
 	o.rwMutex.RLock()
 	defer o.rwMutex.RUnlock()
-	return o.specPbGz, o.specPbGzETag, o.lastModified
+	specPb, etag, err := o.protoCache.Get()
+	if err != nil {
+		return nil, "", time.Time{}, err
+	}
+	return specPb, etag, o.lastModified, nil
 }
 
 func (o *OpenAPIService) UpdateSpec(openapiSpec *spec.Swagger) (err error) {
-	specBytes, err := jsoniter.ConfigCompatibleWithStandardLibrary.Marshal(openapiSpec)
-	if err != nil {
-		return err
+	o.rwMutex.Lock()
+	defer o.rwMutex.Unlock()
+	o.jsonCache = cache{
+		BuildCache: func() ([]byte, error) {
+			return jsoniter.ConfigCompatibleWithStandardLibrary.Marshal(openapiSpec)
+		},
 	}
-	specPb, err := ToProtoBinary(specBytes)
-	if err != nil {
-		return err
+	o.protoCache = cache{
+		BuildCache: func() ([]byte, error) {
+			json, _, err := o.jsonCache.Get()
+			if err != nil {
+				return nil, err
+			}
+			return ToProtoBinary(json)
+		},
 	}
-	specPbGz := toGzip(specPb)
-
-	specBytesETag := computeETag(specBytes)
-	specPbETag := computeETag(specPb)
-	specPbGzETag := computeETag(specPbGz)
-
 	lastModified := time.Now()
-
-	o.rwMutex.Lock()
-	defer o.rwMutex.Unlock()
-
-	o.specBytes = specBytes
-	o.specPb = specPb
-	o.specPbGz = specPbGz
-	o.specBytesETag = specBytesETag
-	o.specPbETag = specPbETag
-	o.specPbGzETag = specPbGzETag
 	o.lastModified = lastModified
 
 	return nil
@@ -206,7 +214,7 @@ func (o *OpenAPIService) RegisterOpenAPIVersionedService(servePath string, handl
 	accepted := []struct {
 		Type           string
 		SubType        string
-		GetDataAndETag func() ([]byte, string, time.Time)
+		GetDataAndETag func() ([]byte, string, time.Time, error)
 	}{
 		{"application", "json", o.getSwaggerBytes},
 		{"application", "[email protected]+protobuf", o.getSwaggerPbBytes},
@@ -230,7 +238,12 @@ func (o *OpenAPIService) RegisterOpenAPIVersionedService(servePath string, handl
 					}
 
 					// serve the first matching media type in the sorted clause list
-					data, etag, lastModified := accepts.GetDataAndETag()
+					data, etag, lastModified, err := accepts.GetDataAndETag()
+					if err != nil {
+						klog.Errorf("Error in OpenAPI handler: %s", err)
+						w.WriteHeader(http.StatusInternalServerError)
+						return
+					}
 					w.Header().Set("Etag", etag)
 					// ServeContent will take care of caching using eTag.
 					http.ServeContent(w, r, servePath, lastModified, bytes.NewReader(data))

pkg/handler/handler_test.go

@@ -1,6 +1,7 @@
 package handler
 
 import (
+	"errors"
 	"io/ioutil"
 	"math"
 	"net/http"
@@ -177,3 +178,34 @@ func TestToProtoBinary(t *testing.T) {
 	}
 	// TODO: add some kind of roundtrip test here
 }
+
+func TestCache(t *testing.T) {
+	calledCount := 0
+	expectedBytes := []byte("ABC")
+	cacheObj := cache{
+		BuildCache: func() ([]byte, error) {
+			calledCount++
+			return expectedBytes, nil
+		},
+	}
+	bytes, _, _ := cacheObj.Get()
+	if string(bytes) != string(expectedBytes) {
+		t.Fatalf("got value of %q from cache (expected %q)", bytes, expectedBytes)
+	}
+	cacheObj.Get()
+	if calledCount != 1 {
+		t.Fatalf("expected BuildCache to be called once (called %d times)", calledCount)
+	}
+}
+
+func TestCacheError(t *testing.T) {
+	cacheObj := cache{
+		BuildCache: func() ([]byte, error) {
+			return nil, errors.New("cache error")
+		},
+	}
+	_, _, err := cacheObj.Get()
+	if err == nil {
+		t.Fatalf("expected non-nil err from cache.Get()")
+	}
+}