feat: support redisv9 (#30)

* feat: support redisv9

* feat: add test for redis

* feat: fix muzzle

.github/workflows/muzzle.yml

@@ -25,4 +25,4 @@ jobs:
       run: make build
 
     - name: Test
-      run: go test -short -timeout 50m -v github.com/alibaba/opentelemetry-go-auto-instrumentation/test -run ".*Muzzle"
+      run: go test -timeout 50m -v github.com/alibaba/opentelemetry-go-auto-instrumentation/test -run ".*Muzzle"

pkg/rules/goredis/goredis_data_type.go

@@ -0,0 +1,12 @@
+//go:build ignore
+
+package goredis
+
+import (
+	redis "github.com/redis/go-redis/v9"
+)
+
+type goRedisRequest struct {
+	cmd      redis.Cmder
+	endpoint string
+}

pkg/rules/goredis/goredis_otel_instrumenter.go

@@ -0,0 +1,44 @@
+//go:build ignore
+
+package goredis
+
+import (
+	"github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/inst-api-semconv/instrumenter/db"
+	"github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/inst-api/instrumenter"
+)
+
+type goRedisAttrsGetter struct {
+}
+
+func (d goRedisAttrsGetter) GetSystem(request goRedisRequest) string {
+	return "redis"
+}
+
+func (d goRedisAttrsGetter) GetUser(request goRedisRequest) string {
+	return ""
+}
+
+func (d goRedisAttrsGetter) GetName(request goRedisRequest) string {
+	// TODO: parse database name from dsn
+	return ""
+}
+
+func (d goRedisAttrsGetter) GetConnectionString(request goRedisRequest) string {
+	return request.endpoint
+}
+
+func (d goRedisAttrsGetter) GetStatement(request goRedisRequest) string {
+	return request.cmd.String()
+}
+
+func (d goRedisAttrsGetter) GetOperation(request goRedisRequest) string {
+	return request.cmd.FullName()
+}
+
+func BuildGoRedisOtelInstrumenter() *instrumenter.Instrumenter[goRedisRequest, interface{}] {
+	builder := instrumenter.Builder[goRedisRequest, interface{}]{}
+	getter := goRedisAttrsGetter{}
+	return builder.Init().SetSpanNameExtractor(&db.DBSpanNameExtractor[goRedisRequest]{Getter: getter}).SetSpanKindExtractor(&instrumenter.AlwaysClientExtractor[goRedisRequest]{}).
+		AddAttributesExtractor(&db.DbClientAttrsExtractor[goRedisRequest, any, goRedisAttrsGetter]{Base: db.DbClientCommonAttrsExtractor[goRedisRequest, any, goRedisAttrsGetter]{Getter: getter}}).
+		BuildInstrumenter()
+}

pkg/rules/goredis/rule.go

@@ -0,0 +1,34 @@
+package goredis
+
+import "github.com/alibaba/opentelemetry-go-auto-instrumentation/api"
+
+func init() {
+	api.NewRule("github.com/redis/go-redis/v9", "NewClient", "", "", "afterNewRedisClient").
+		WithFileDeps("goredis_data_type.go", "goredis_otel_instrumenter.go").
+		Register()
+
+	api.NewRule("github.com/redis/go-redis/v9", "NewFailoverClient", "", "", "afterNewFailOverRedisClient").
+		WithVersion("[9.0.5,9.5.2)").
+		WithFileDeps("goredis_data_type.go", "goredis_otel_instrumenter.go").
+		Register()
+
+	api.NewRule("github.com/redis/go-redis/v9", "NewSentinelClient", "", "", "afterNewSentinelClient").
+		WithVersion("[9.0.5,9.5.2)").
+		WithFileDeps("goredis_data_type.go", "goredis_otel_instrumenter.go").
+		Register()
+
+	api.NewRule("github.com/redis/go-redis/v9", "Conn", "*Client", "", "afterClientConn").
+		WithVersion("[9.0.5,9.5.2)").
+		WithFileDeps("goredis_data_type.go", "goredis_otel_instrumenter.go").
+		Register()
+
+	api.NewRule("github.com/redis/go-redis/v9", "NewClusterClient", "", "", "afterNewClusterClient").
+		WithVersion("[9.0.5,9.5.2)").
+		WithFileDeps("goredis_data_type.go", "goredis_otel_instrumenter.go").
+		Register()
+
+	api.NewRule("github.com/redis/go-redis/v9", "NewRing", "", "", "afterNewRingClient").
+		WithVersion("[9.0.5,9.5.2)").
+		WithFileDeps("goredis_data_type.go", "goredis_otel_instrumenter.go").
+		Register()
+}

pkg/rules/goredis/setup.go

@@ -0,0 +1,107 @@
+//go:build ignore
+
+package goredis
+
+import (
+	"context"
+	redis "github.com/redis/go-redis/v9"
+	"net"
+	"strings"
+)
+
+var goRedisInstrumenter = BuildGoRedisOtelInstrumenter()
+
+func afterNewRedisClient(call redis.CallContext, client *redis.Client) {
+	client.AddHook(newOtRedisHook(client.Options().Addr))
+}
+
+func afterNewFailOverRedisClient(call redis.CallContext, client *redis.Client) {
+	client.AddHook(newOtRedisHook(client.Options().Addr))
+}
+
+func afterNewClusterClient(call redis.CallContext, client *redis.ClusterClient) {
+	client.OnNewNode(func(rdb *redis.Client) {
+		rdb.AddHook(newOtRedisHook(rdb.Options().Addr))
+	})
+}
+
+func afterNewRingClient(call redis.CallContext, client *redis.Ring) {
+	client.OnNewNode(func(rdb *redis.Client) {
+		rdb.AddHook(newOtRedisHook(rdb.Options().Addr))
+	})
+}
+
+func afterNewSentinelClient(call redis.CallContext, client *redis.SentinelClient) {
+	client.AddHook(newOtRedisHook(client.String()))
+}
+
+func afterClientConn(call redis.CallContext, client *redis.Conn) {
+	client.AddHook(newOtRedisHook(client.String()))
+}
+
+type otRedisHook struct {
+	Addr string
+}
+
+func newOtRedisHook(addr string) *otRedisHook {
+	return &otRedisHook{
+		Addr: addr,
+	}
+}
+
+func (o *otRedisHook) DialHook(next redis.DialHook) redis.DialHook {
+	return func(ctx context.Context, network, addr string) (net.Conn, error) {
+		conn, err := next(ctx, network, addr)
+		if err != nil {
+			return nil, err
+		}
+		return conn, err
+	}
+}
+
+func (o *otRedisHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook {
+	return func(ctx context.Context, cmd redis.Cmder) error {
+		if strings.Contains(cmd.FullName(), "ping") || strings.Contains(cmd.FullName(), "PING") {
+			return next(ctx, cmd)
+		}
+		request := goRedisRequest{
+			cmd:      cmd,
+			endpoint: o.Addr,
+		}
+		ctx = goRedisInstrumenter.Start(ctx, request)
+		if err := next(ctx, cmd); err != nil {
+			goRedisInstrumenter.End(ctx, request, nil, err)
+			return err
+		}
+		goRedisInstrumenter.End(ctx, request, nil, nil)
+		return nil
+	}
+}
+
+func (o *otRedisHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook {
+	return func(ctx context.Context, cmds []redis.Cmder) error {
+		summary := ""
+		summaryCmds := cmds
+		if len(summaryCmds) > 10 {
+			summaryCmds = summaryCmds[:10]
+		}
+		for i := range summaryCmds {
+			summary += summaryCmds[i].FullName() + "/"
+		}
+		if len(cmds) > 10 {
+			summary += "..."
+		}
+		cmd := redis.NewCmd(ctx, "pipeline", summary)
+		request := goRedisRequest{
+			cmd:      cmd,
+			endpoint: o.Addr,
+		}
+		ctx = goRedisInstrumenter.Start(ctx, request)
+		if err := next(ctx, cmds); err != nil {
+			goRedisInstrumenter.End(ctx, request, nil, err)
+			return err
+		}
+		goRedisInstrumenter.End(ctx, request, nil, nil)
+		return nil
+	}
+}

rule_enabler.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	_ "github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/rules/databasesql"
+	_ "github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/rules/goredis"
 	_ "github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/rules/mongo"
 	_ "github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/rules/test"
 )

test/databasesql_tests.go

@@ -7,6 +7,7 @@ import (
 	"github.com/testcontainers/testcontainers-go/modules/mysql"
 	"log"
 	"testing"
+	"time"
 )
 
 func init() {
@@ -160,6 +161,7 @@ func init8xMySqlContainer() (testcontainers.Container, nat.Port) {
 	if err := mysqlContainer.Start(ctx); err != nil {
 		panic(err)
 	}
+	time.Sleep(5 * time.Second)
 	port, err := mysqlContainer.MappedPort(ctx, "3306")
 	if err != nil {
 		panic(err)

test/infra.go

@@ -95,9 +95,13 @@ func RunInstrument(t *testing.T, args ...string) {
 	}
 }
 
-func TBuildAppNoop(t *testing.T, appName string) {
+func TBuildAppNoop(t *testing.T, appName, mainClass string) {
 	UseApp(appName)
-	RunInstrument(t)
+	if mainClass == "" {
+		RunInstrument(t)
+	} else {
+		RunInstrument(t, "--", mainClass)
+	}
 }
 
 func BuildApp(t *testing.B, appName string) {
@@ -238,7 +242,7 @@ func ExpectContainsNothing(t *testing.T, actualItems []string) {
 	}
 }
 
-func ExecMuzzle(t *testing.T, dependencyName, moduleName string, minVersion, maxVersion *version.Version) {
+func ExecMuzzle(t *testing.T, dependencyName, moduleName string, minVersion, maxVersion *version.Version, mainClass string) {
 	if testing.Short() {
 		t.Skip()
 		return
@@ -268,14 +272,14 @@ func ExecMuzzle(t *testing.T, dependencyName, moduleName string, minVersion, max
 				t.Logf("testing on version %v\n", version.Original())
 				UseApp(moduleName + "/" + testVersion.Original())
 				FetchVersion(t, dependencyName, version.Original())
-				TBuildAppNoop(t, moduleName+"/"+testVersion.Original())
+				TBuildAppNoop(t, moduleName+"/"+testVersion.Original(), mainClass)
 				break
 			}
 		}
 	}
 }
 
-func ExecLatestTest(t *testing.T, dependencyName, moduleName string, minVersion, maxVersion *version.Version, testFunc func(*testing.T, *version.Version, ...string)) {
+func ExecLatestTest(t *testing.T, dependencyName, moduleName string, minVersion, maxVersion *version.Version, testFunc func(*testing.T, ...string)) {
 	if testing.Short() {
 		t.Skip()
 		return
@@ -302,5 +306,5 @@ func ExecLatestTest(t *testing.T, dependencyName, moduleName string, minVersion,
 	latestTestVersion := testVersions[len(testVersions)-1]
 	UseApp(moduleName + "/" + latestTestVersion.Original())
 	FetchVersion(t, dependencyName, latestVersion.Original())
-	testFunc(t, latestTestVersion)
+	testFunc(t)
 }

test/integration_test.go

@@ -7,8 +7,8 @@ import (
 
 func TestPlugins(t *testing.T) {
 	for _, c := range TestCases {
-		if c == nil {
-			t.Skip()
+		if c == nil || c.TestName != "redis-9.0.5-universal-test" {
+			continue
 		}
 		if c.IsMuzzleCheck || c.IsLatestDepthCheck {
 			continue
@@ -22,21 +22,21 @@ func TestPlugins(t *testing.T) {
 func TestMuzzle(t *testing.T) {
 	for _, c := range TestCases {
 		if c == nil {
-			t.Skip()
+			continue
 		}
 		if !c.IsMuzzleCheck {
 			continue
 		}
 		t.Run(c.TestName, func(t *testing.T) {
-			ExecMuzzle(t, c.DependencyName, c.ModuleName, c.MinVersion, c.MaxVersion)
+			ExecMuzzle(t, c.DependencyName, c.ModuleName, c.MinVersion, c.MaxVersion, c.MuzzleMainClass)
 		})
 	}
 }
 
 func TestLatest(t *testing.T) {
 	for _, c := range TestCases {
 		if c == nil {
-			t.Skip()
+			continue
 		}
 		if !c.IsLatestDepthCheck {
 			continue

test/mongo_tests.go

@@ -2,37 +2,29 @@ package test
 
 import (
 	"context"
-	"github.com/alibaba/opentelemetry-go-auto-instrumentation/test/version"
 	"github.com/docker/go-connections/nat"
 	"github.com/testcontainers/testcontainers-go"
 	"log"
 	"testing"
+	"time"
 )
 
-const dependency_name = "go.mongodb.org/mongo-driver"
-const module_name = "mongo"
+const mongo_dependency_name = "go.mongodb.org/mongo-driver"
+const mongo_module_name = "mongo"
 
 func init() {
-	TestCases = append(TestCases, NewGeneralTestCase("mongo-1.11.1-test", dependency_name, module_name, "v1.11.1", "v1.15.1", "1.18", "", TestGoMongo111),
-		NewMuzzleTestCase("mongo-1.11.1-muzzle", dependency_name, module_name, "v1.11.1", "v1.15.1", "1.18", ""),
-		NewLatestDepthTestCase("mongo-1.11.1-latestDepth", dependency_name, module_name, "v1.11.1", "v1.15.1", "1.18", "", TestMongoLatest))
+	TestCases = append(TestCases, NewGeneralTestCase("mongo-1.11.1-test", mongo_dependency_name, mongo_module_name, "v1.11.1", "v1.15.1", "1.18", "", TestCrudMongo),
+		NewMuzzleTestCase("mongo-1.11.1-muzzle", mongo_dependency_name, mongo_module_name, "v1.11.1", "v1.15.1", "1.18", "", ""),
+		NewLatestDepthTestCase("mongo-1.11.1-latestDepth", mongo_dependency_name, mongo_module_name, "v1.11.1", "v1.15.1", "1.18", "", TestCrudMongo))
 }
 
-func TestGoMongo111(t *testing.T, env ...string) {
+func TestCrudMongo(t *testing.T, env ...string) {
 	mongoC, mongoPort := initMongoContainer()
 	defer clearMongoContainer(mongoC)
 	UseApp("mongo/v1.11.1")
-	RunInstrument(t, "-debuglog")
+	RunInstrument(t, "-debuglog", "--", "test_crud_mongo.go")
 	env = append(env, "MONGO_PORT="+mongoPort.Port())
-	RunApp(t, "v1.11.1", env...)
-}
-
-func TestMongoLatest(t *testing.T, v *version.Version, env ...string) {
-	mongoC, mongoPort := initMongoContainer()
-	defer clearMongoContainer(mongoC)
-	RunInstrument(t, "-debuglog")
-	env = append(env, "MONGO_PORT="+mongoPort.Port())
-	RunApp(t, v.Original(), env...)
+	RunApp(t, "test_crud_mongo", env...)
 }
 
 func initMongoContainer() (testcontainers.Container, nat.Port) {
@@ -47,6 +39,7 @@ func initMongoContainer() (testcontainers.Container, nat.Port) {
 	if err != nil {
 		panic(err)
 	}
+	time.Sleep(5 * time.Second)
 	port, err := mongoC.MappedPort(context.Background(), "27017")
 	if err != nil {
 		panic(err)

test/redis/v9.0.5/go.mod

@@ -0,0 +1,11 @@
+module redis/v9.0.5
+
+go 1.21
+
+replace github.com/alibaba/opentelemetry-go-auto-instrumentation => ../../../../opentelemetry-go-auto-instrumentation
+
+require (
+	github.com/alibaba/opentelemetry-go-auto-instrumentation v0.0.0-00010101000000-000000000000
+	github.com/redis/go-redis/v9 v9.0.5
+	go.opentelemetry.io/otel/sdk v1.28.0
+)