github.com/erda-project/erda-infra@v1.0.10-0.20240327085753-f3a249292aeb/pkg/trace/inject/redis/client.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package redis
    16  
    17  import (
    18  	"github.com/go-redis/redis"
    19  
    20  	"github.com/erda-project/erda-infra/pkg/trace/inject/hook"
    21  )
    22  
    23  // Client .
    24  type Client interface {
    25  	WrapProcess(fn func(oldProcess func(redis.Cmder) error) func(redis.Cmder) error)
    26  	WrapProcessPipeline(fn func(oldProcess func([]redis.Cmder) error) func([]redis.Cmder) error)
    27  }
    28  
    29  // Wrap .
    30  func Wrap(client Client, opts ...Option) {
    31  	if client == nil {
    32  		return
    33  	}
    34  	cfg := newConfig("redis", opts...)
    35  	client.WrapProcess(newProcessWrapper(cfg))
    36  	client.WrapProcessPipeline(newProcessPipeline(cfg))
    37  }
    38  
    39  //go:noinline
    40  func originalNewClient(opts *redis.Options) *redis.Client {
    41  	return redis.NewClient(opts)
    42  }
    43  
    44  //go:noinline
    45  func originalNewFailoverClient(opts *redis.FailoverOptions) *redis.Client {
    46  	return redis.NewFailoverClient(opts)
    47  }
    48  
    49  //go:noinline
    50  func originalNewClusterClient(opts *redis.ClusterOptions) *redis.ClusterClient {
    51  	return redis.NewClusterClient(opts)
    52  }
    53  
    54  //go:noinline
    55  func wrappedNewClient(opts *redis.Options) *redis.Client {
    56  	client := originalNewClient(opts)
    57  	Wrap(client)
    58  	return client
    59  }
    60  
    61  //go:noinline
    62  func wrappedNewFailoverClient(opts *redis.FailoverOptions) *redis.Client {
    63  	client := originalNewFailoverClient(opts)
    64  	Wrap(client)
    65  	return client
    66  }
    67  
    68  //go:noinline
    69  func wrappedNewClusterClient(opts *redis.ClusterOptions) *redis.ClusterClient {
    70  	client := originalNewClusterClient(opts)
    71  	Wrap(client)
    72  	return client
    73  }
    74  
    75  func init() {
    76  	hook.Hook(redis.NewClient, wrappedNewClient, originalNewClient)
    77  	hook.Hook(redis.NewFailoverClient, wrappedNewFailoverClient, originalNewFailoverClient)
    78  	hook.Hook(redis.NewClusterClient, wrappedNewClusterClient, originalNewClusterClient)
    79  }