[BEAM-14347] Add some benchmarks for generic registration

sdks/go/pkg/beam/register/register_test.go

@@ -20,9 +20,12 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/exec"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx/schema"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx"
@@ -672,3 +675,133 @@ func (fn *PartialCombiner2) AddInput(i int, c CustomType) int {
 func (fn *PartialCombiner2) MergeAccumulators(i1 int, i2 int) int {
 	return i1 + i2
 }
+
+// Foo is a struct with a method for measuring method invocation
+// overhead for StructuralDoFns.
+type Foo struct {
+	A int
+}
+
+// ProcessElement is a method for measuring a baseline of structural dofn overhead.
+func (f *Foo) ProcessElement(b CustomType) int {
+	return f.A + b.val
+}
+
+func MakeMultiEdge(f *graph.DoFn) graph.MultiEdge {
+	return graph.MultiEdge{
+		DoFn: f,
+	}
+}
+
+type callerCustomTypeГint struct {
+	fn func(CustomType) int
+}
+
+func funcMakerCustomTypeГint(fn interface{}) reflectx.Func {
+	f := fn.(func(CustomType) int)
+	return &callerCustomTypeГint{fn: f}
+}
+
+func (c *callerCustomTypeГint) Name() string {
+	return reflectx.FunctionName(c.fn)
+}
+
+func (c *callerCustomTypeГint) Type() reflect.Type {
+	return reflect.TypeOf(c.fn)
+}
+
+func (c *callerCustomTypeГint) Call(args []interface{}) []interface{} {
+	out0 := c.fn(args[0].(CustomType))
+	return []interface{}{out0}
+}
+
+func (c *callerCustomTypeГint) Call1x1(arg0 interface{}) interface{} {
+	return c.fn(arg0.(CustomType))
+}
+
+func wrapMakerFoo(fn interface{}) map[string]reflectx.Func {
+	dfn := fn.(*Foo)
+	return map[string]reflectx.Func{
+		"ProcessElement": reflectx.MakeFunc(func(a0 CustomType) int { return dfn.ProcessElement(a0) }),
+	}
+}
+
+func GeneratedOptimizationCalls() {
+	runtime.RegisterType(reflect.TypeOf((*Foo)(nil)).Elem())
+	schema.RegisterType(reflect.TypeOf((*Foo)(nil)).Elem())
+	runtime.RegisterType(reflect.TypeOf((*CustomType)(nil)).Elem())
+	schema.RegisterType(reflect.TypeOf((*CustomType)(nil)).Elem())
+	reflectx.RegisterFunc(reflect.TypeOf((*func(CustomType) int)(nil)).Elem(), funcMakerCustomTypeГint)
+	reflectx.RegisterStructWrapper(reflect.TypeOf((*Foo)(nil)).Elem(), wrapMakerFoo)
+}
+
+// BenchmarkMethodCalls measures the overhead of invoking several different methods after performing
+// different types of registration. The unoptimized calls don't perform any optimization. The
+// GenericRegistration calls first register the DoFn being used with this package's generic registration
+// functions. This is the preferred path for users. The GeneratedShims calls call various registration
+// functions, mirroring the behavior of the shims generated by the code generator. This is not the
+// recommended path for most users - if these are materially better than the generic benchmarks,
+// this package requires further optimization.
+//
+// BenchmarkMethodCalls/MakeFunc_Unoptimized-16                      11480814	        88.35 ns/op
+// BenchmarkMethodCalls/NewFn_Unoptimized-16                         875199	      		1385 ns/op
+// BenchmarkMethodCalls/EncodeMultiEdge_Unoptimized-16               1000000	      	1063 ns/op
+//
+// BenchmarkMethodCalls/MakeFunc_GenericRegistration-16             16266259	        72.07 ns/op
+// BenchmarkMethodCalls/NewFn_GenericRegistration-16                 1000000	      	1108 ns/op
+// BenchmarkMethodCalls/EncodeMultiEdge_GenericRegistration-16       1000000	      	1052 ns/op
+//
+// BenchmarkMethodCalls/MakeFunc_GeneratedShims#01-16               16400914	        69.17 ns/op
+// BenchmarkMethodCalls/NewFn_GeneratedShims#01-16                   1000000	      	1099 ns/op
+// BenchmarkMethodCalls/EncodeMultiEdge_GeneratedShims#01-16         1000000	      	1071 ns/op
+func BenchmarkMethodCalls(b *testing.B) {
+	f := &Foo{A: 3}
+	g, err := graph.NewFn(&Foo{A: 5})
+	if err != nil {
+		panic(err)
+	}
+	gDoFn, err := graph.AsDoFn(g, 1)
+	if err != nil {
+		panic(err)
+	}
+	me := MakeMultiEdge(gDoFn)
+
+	var aFunc reflectx.Func
+	var aFn *graph.Fn
+	var aME interface{}
+
+	// We need to do this registration just to get it to not panic when encoding the multi-edge with no additional optimization.
+	// This is currently required of users anyways
+	runtime.RegisterType(reflect.TypeOf((*Foo)(nil)))
+	tests := []struct {
+		name         string
+		fn           func()
+		registration func()
+	}{
+		// No optimization performed at all
+		{"MakeFunc_Unoptimized", func() { aFunc = reflectx.MakeFunc(f.ProcessElement) }, func() { /*No op*/ }}, // Used in graph deserialization
+		{"NewFn_Unoptimized", func() { aFn, _ = graph.NewFn(f) }, func() { /*No op*/ }},                        // Used in graph construction (less valuable)
+		{"EncodeMultiEdge_Unoptimized", func() { aME, _ = graphx.EncodeMultiEdge(&me) }, func() { /*No op*/ }}, // Used in graph serialization at execution time
+
+		// Perform some generic registration to optimize execution
+		{"MakeFunc_GenericRegistration", func() { aFunc = reflectx.MakeFunc(f.ProcessElement) }, func() { DoFn1x1[CustomType, int](f) }}, // Used in graph deserialization
+		{"NewFn_GenericRegistration", func() { aFn, _ = graph.NewFn(f) }, func() { DoFn1x1[CustomType, int](f) }},                        // Used in graph construction (less valuable)
+		{"EncodeMultiEdge_GenericRegistration", func() { aME, _ = graphx.EncodeMultiEdge(&me) }, func() { DoFn1x1[CustomType, int](f) }}, // Used in graph serialization at execution time
+
+		// Perform some registration via copies of the code generator's shims
+		{"MakeFunc_GeneratedShims", func() { aFunc = reflectx.MakeFunc(f.ProcessElement) }, func() { GeneratedOptimizationCalls() }},   // Used in graph deserialization
+		{"NewFn_GeneratedShims", func() { aFn, _ = graph.NewFn(f) }, func() { GeneratedOptimizationCalls() }},                          // Used in graph construction (less valuable)
+		{"EncodeMultiEdge_GeneratedShims", func() { aME, err = graphx.EncodeMultiEdge(&me) }, func() { GeneratedOptimizationCalls() }}, // Used in graph serialization at execution time
+	}
+	for _, test := range tests {
+		test.registration()
+		b.Run(test.name, func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				test.fn()
+			}
+		})
+	}
+	b.Log(aFunc)
+	b.Log(aFn)
+	b.Log(aME)
+}