github.com/jtzjtz/kit@v1.0.2/conn/mongo_pool/mongo_pool.go (about)

     1  package mongo_pool
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"go.mongodb.org/mongo-driver/mongo"
     7  	"go.mongodb.org/mongo-driver/mongo/options"
     8  	"log"
     9  	"sync"
    10  	"time"
    11  )
    12  
    13  const (
    14  	INITIAL_CONNECTION = 5
    15  	MAX_CONNECTION     = 50
    16  	AVAILABLE          = false
    17  	USED               = true
    18  )
    19  
    20  var mu sync.RWMutex
    21  
    22  type mongodata struct {
    23  	client *mongo.Client
    24  	pos    int
    25  	flag   bool
    26  }
    27  
    28  type ClientPool struct {
    29  	uri        string
    30  	clientList [MAX_CONNECTION]mongodata
    31  	size       int
    32  }
    33  
    34  func NewClientPool(uri string) (*ClientPool, error) {
    35  	if len(uri) < 1 {
    36  		return nil, errors.New("connect uri is empty")
    37  	}
    38  	cp := &ClientPool{
    39  		uri:        uri,
    40  		clientList: [50]mongodata{},
    41  		size:       0,
    42  	}
    43  	for size := 0; size < INITIAL_CONNECTION || size < MAX_CONNECTION; size++ {
    44  		err := cp.allocateCToPool(size)
    45  		log.Fatal("init - initial create the connect conn failed, size:", size, err)
    46  		return nil, err
    47  	}
    48  	return cp, nil
    49  }
    50  
    51  func (cp *ClientPool) Dbconnect() (client *mongo.Client, err error) {
    52  	client, err = mongo.NewClient(options.Client().ApplyURI(cp.uri))
    53  	if err != nil {
    54  		log.Fatal("Dbconnect - connect mongodb failed", err)
    55  		return nil, err
    56  	}
    57  	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
    58  	defer cancel()
    59  	err = client.Connect(ctx)
    60  	if err != nil {
    61  		log.Fatal("Dbconnect - connect mongodb ctx failed", err)
    62  		return nil, err
    63  	}
    64  
    65  	return client, nil
    66  }
    67  
    68  func (cp *ClientPool) Dbdisconnect(client *mongo.Client) (err error) {
    69  	err = client.Disconnect(context.TODO())
    70  	if err != nil {
    71  		log.Fatal("Dbdisconnect - disconnect mongodb failed", err)
    72  	}
    73  	return err
    74  }
    75  
    76  func (cp *ClientPool) allocateCToPool(pos int) (err error) {
    77  	cp.clientList[pos].client, err = cp.Dbconnect()
    78  	if err != nil {
    79  		log.Fatal("allocateCToPool - allocateCToPool failed,position: ", pos, err)
    80  		return err
    81  	}
    82  
    83  	cp.clientList[pos].flag = USED
    84  	cp.clientList[pos].pos = pos
    85  	return nil
    86  }
    87  
    88  func (cp *ClientPool) getCToPool(pos int) {
    89  	cp.clientList[pos].flag = USED
    90  }
    91  
    92  func (cp *ClientPool) putCBackPool(pos int) {
    93  	cp.clientList[pos].flag = AVAILABLE
    94  }
    95  
    96  func (cp *ClientPool) GetClient() (mongoclient *mongodata, err error) {
    97  	mu.RLock()
    98  	for i := 1; i < cp.size; i++ {
    99  		if cp.clientList[i].flag == AVAILABLE {
   100  			return &cp.clientList[i], nil
   101  		}
   102  	}
   103  	mu.RUnlock()
   104  	mu.Lock()
   105  	defer mu.Unlock()
   106  	if cp.size < MAX_CONNECTION {
   107  		err = cp.allocateCToPool(cp.size)
   108  		if err != nil {
   109  			log.Fatal("GetClient - DB pooling allocate failed", err)
   110  			return nil, err
   111  		}
   112  		pos := cp.size
   113  		cp.size++
   114  		return &cp.clientList[pos], nil
   115  	} else {
   116  		log.Fatal("GetClient - DB pooling is fulled")
   117  		return nil, errors.New("DB pooling is fulled")
   118  	}
   119  }
   120  
   121  func (cp *ClientPool) ReleaseClient(mongoclient *mongodata) {
   122  	mu.Lock()
   123  	cp.putCBackPool(mongoclient.pos)
   124  	mu.Unlock()
   125  }