tokenizer

ai-chat-service/services/grpc-client/default.go

@@ -0,0 +1,28 @@
+package grpc_client
+
+import (
+	"ai-chat-service/pkg/log"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+type ServiceClient interface {
+	GetPool(addr string) ClientPool
+}
+type DefaultClient struct {
+}
+
+func (c *DefaultClient) GetPool(addr string) ClientPool {
+	pool, err := NewPool(addr, c.getOptions()...)
+	if err != nil {
+		log.Error(err)
+		return nil
+	}
+	return pool
+}
+
+func (c *DefaultClient) getOptions() []grpc.DialOption {
+	opts := make([]grpc.DialOption, 0)
+	opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	return opts
+}

ai-chat-service/services/grpc-client/grpc_client_pool.go

@@ -0,0 +1,48 @@
+package grpc_client
+
+import (
+	"ai-chat-service/pkg/log"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
+	"sync"
+)
+
+type ClientPool interface {
+	Get() *grpc.ClientConn
+	Put(*grpc.ClientConn)
+}
+
+type clientPool struct {
+	pool sync.Pool
+}
+
+func NewPool(target string, opts ...grpc.DialOption) (ClientPool, error) {
+	return &clientPool{
+		pool: sync.Pool{
+			New: func() any {
+				conn, err := grpc.NewClient(target, opts...)
+				if err != nil {
+					log.Error(err)
+					return nil
+				}
+				return conn
+			},
+		},
+	}, nil
+}
+
+func (c *clientPool) Get() *grpc.ClientConn {
+	conn := c.pool.Get().(*grpc.ClientConn)
+	if conn.GetState() == connectivity.Shutdown || conn.GetState() == connectivity.TransientFailure {
+		conn.Close()
+		conn = c.pool.New().(*grpc.ClientConn)
+	}
+	return conn
+}
+func (c *clientPool) Put(conn *grpc.ClientConn) {
+	if conn.GetState() == connectivity.Shutdown || conn.GetState() == connectivity.TransientFailure {
+		conn.Close()
+		return
+	}
+	c.pool.Put(conn)
+}