Better cancellation (#148)

Don’t override global HTTP handler 😉

hack/generator/cmd/gen/gen.go

@@ -7,7 +7,6 @@ package gen
 
 import (
 	"context"
-	"net/http"
 
 	"github.com/Azure/k8s-infra/hack/generator/pkg/codegen"
 	"github.com/Azure/k8s-infra/hack/generator/pkg/xcobra"
@@ -16,32 +15,13 @@ import (
 	"k8s.io/klog/v2"
 )
 
-type cancellableTransport struct {
-	ctx   context.Context
-	inner http.RoundTripper
-}
-
-var _ http.RoundTripper = &cancellableTransport{} // interface assertion
-
-func (transport *cancellableTransport) RoundTrip(request *http.Request) (*http.Response, error) {
-	if transport.ctx.Err() != nil { // check for cancellation
-		return nil, transport.ctx.Err()
-	}
-
-	return transport.inner.RoundTrip(request)
-}
-
 // NewGenCommand creates a new cobra Command when invoked from the command line
 func NewGenCommand() (*cobra.Command, error) {
 	cmd := &cobra.Command{
 		Use:   "gen",
 		Short: "generate K8s infrastructure resources from Azure deployment template schema",
 		Run: xcobra.RunWithCtx(func(ctx context.Context, cmd *cobra.Command, args []string) error {
 
-			// HACK HACK HACK: make all HTTP clients cancellable as gojsonschema doesn't permit
-			// either cancellation or overriding the HTTP fetch
-			http.DefaultTransport = &cancellableTransport{ctx, http.DefaultTransport}
-
 			cg, err := codegen.NewCodeGenerator("azure-cloud.yaml")
 			if err != nil {
 				klog.Errorf("Error creating code generator: %v\n", err)

hack/generator/go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.6.1
 	github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb // indirect
+	github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415
 	github.com/xeipuuv/gojsonschema v1.2.0
 	golang.org/x/net v0.0.0-20191209160850-c0dbc17a3553 // indirect
 	golang.org/x/sys v0.0.0-20191228213918-04cbcbbfeed8 // indirect

hack/generator/pkg/codegen/code_generator.go

@@ -21,6 +21,7 @@ import (
 	"github.com/Azure/k8s-infra/hack/generator/pkg/astmodel"
 	"github.com/Azure/k8s-infra/hack/generator/pkg/config"
 	"github.com/Azure/k8s-infra/hack/generator/pkg/jsonast"
+	"github.com/xeipuuv/gojsonreference"
 	"github.com/xeipuuv/gojsonschema"
 	"gopkg.in/yaml.v3"
 
@@ -300,11 +301,55 @@ func (fs *cancellableFileSystem) Open(source string) (http.File, error) {
 	return os.Open(source)
 }
 
+type cancellableJSONLoaderFactory struct {
+	ctx   context.Context
+	inner gojsonschema.JSONLoaderFactory
+}
+
+var _ gojsonschema.JSONLoaderFactory = &cancellableJSONLoaderFactory{}
+
+func (factory *cancellableJSONLoaderFactory) New(source string) gojsonschema.JSONLoader {
+	return &cancellableJSONLoader{factory.ctx, factory.inner.New(source)}
+}
+
+type cancellableJSONLoader struct {
+	ctx   context.Context
+	inner gojsonschema.JSONLoader
+}
+
+var _ gojsonschema.JSONLoader = &cancellableJSONLoader{}
+
+func (loader *cancellableJSONLoader) LoadJSON() (interface{}, error) {
+	if loader.ctx.Err() != nil { // check for cancellation
+		return nil, loader.ctx.Err()
+	}
+
+	return loader.inner.LoadJSON()
+}
+
+func (loader *cancellableJSONLoader) JsonSource() interface{} {
+	return loader.inner.JsonSource()
+}
+
+func (loader *cancellableJSONLoader) JsonReference() (gojsonreference.JsonReference, error) {
+	if loader.ctx.Err() != nil { // check for cancellation
+		return gojsonreference.JsonReference{}, loader.ctx.Err()
+	}
+
+	return loader.inner.JsonReference()
+}
+
+func (loader *cancellableJSONLoader) LoaderFactory() gojsonschema.JSONLoaderFactory {
+	return &cancellableJSONLoaderFactory{loader.ctx, loader.inner.LoaderFactory()}
+}
+
 func loadSchema(ctx context.Context, source string) (*gojsonschema.Schema, error) {
 	sl := gojsonschema.NewSchemaLoader()
-	// note that we "configure" the DefaultClient in gen.go to cancel HTTP calls
-	// the cancellableFS here only handles actual FS calls
-	loader := gojsonschema.NewReferenceLoaderFileSystem(source, &cancellableFileSystem{ctx})
+	loader := &cancellableJSONLoader{
+		ctx,
+		gojsonschema.NewReferenceLoaderFileSystem(source, &cancellableFileSystem{ctx}),
+	}
+
 	schema, err := sl.Compile(loader)
 	if err != nil {
 		return nil, fmt.Errorf("error loading schema from '%v' (%w)", source, err)