github.com/lingyao2333/mo-zero@v1.4.1/core/stores/mon/clientmanager.go (about)

     1  package mon
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"time"
     7  
     8  	"github.com/lingyao2333/mo-zero/core/syncx"
     9  	"go.mongodb.org/mongo-driver/mongo"
    10  	mopt "go.mongodb.org/mongo-driver/mongo/options"
    11  )
    12  
    13  const defaultTimeout = time.Second
    14  
    15  var clientManager = syncx.NewResourceManager()
    16  
    17  // ClosableClient wraps *mongo.Client and provides a Close method.
    18  type ClosableClient struct {
    19  	*mongo.Client
    20  }
    21  
    22  // Close disconnects the underlying *mongo.Client.
    23  func (cs *ClosableClient) Close() error {
    24  	return cs.Client.Disconnect(context.Background())
    25  }
    26  
    27  // Inject injects a *mongo.Client into the client manager.
    28  // Typically, this is used to inject a *mongo.Client for test purpose.
    29  func Inject(key string, client *mongo.Client) {
    30  	clientManager.Inject(key, &ClosableClient{client})
    31  }
    32  
    33  func getClient(url string) (*mongo.Client, error) {
    34  	val, err := clientManager.GetResource(url, func() (io.Closer, error) {
    35  		cli, err := mongo.Connect(context.Background(), mopt.Client().ApplyURI(url))
    36  		if err != nil {
    37  			return nil, err
    38  		}
    39  
    40  		concurrentSess := &ClosableClient{
    41  			Client: cli,
    42  		}
    43  
    44  		return concurrentSess, nil
    45  	})
    46  	if err != nil {
    47  		return nil, err
    48  	}
    49  
    50  	return val.(*ClosableClient).Client, nil
    51  }