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

     1  package mongo
     2  
     3  import (
     4  	"io"
     5  	"time"
     6  
     7  	"github.com/globalsign/mgo"
     8  	"github.com/lingyao2333/mo-zero/core/logx"
     9  	"github.com/lingyao2333/mo-zero/core/syncx"
    10  )
    11  
    12  const (
    13  	defaultConcurrency = 50
    14  	defaultTimeout     = time.Second
    15  )
    16  
    17  var sessionManager = syncx.NewResourceManager()
    18  
    19  type concurrentSession struct {
    20  	*mgo.Session
    21  	limit syncx.TimeoutLimit
    22  }
    23  
    24  func (cs *concurrentSession) Close() error {
    25  	cs.Session.Close()
    26  	return nil
    27  }
    28  
    29  func getConcurrentSession(url string) (*concurrentSession, error) {
    30  	val, err := sessionManager.GetResource(url, func() (io.Closer, error) {
    31  		mgoSession, err := mgo.Dial(url)
    32  		if err != nil {
    33  			return nil, err
    34  		}
    35  
    36  		concurrentSess := &concurrentSession{
    37  			Session: mgoSession,
    38  			limit:   syncx.NewTimeoutLimit(defaultConcurrency),
    39  		}
    40  
    41  		return concurrentSess, nil
    42  	})
    43  	if err != nil {
    44  		return nil, err
    45  	}
    46  
    47  	return val.(*concurrentSession), nil
    48  }
    49  
    50  func (cs *concurrentSession) putSession(session *mgo.Session) {
    51  	if err := cs.limit.Return(); err != nil {
    52  		logx.Error(err)
    53  	}
    54  
    55  	// anyway, we need to close the session
    56  	session.Close()
    57  }
    58  
    59  func (cs *concurrentSession) takeSession(opts ...Option) (*mgo.Session, error) {
    60  	o := defaultOptions()
    61  	for _, opt := range opts {
    62  		opt(o)
    63  	}
    64  
    65  	if err := cs.limit.Borrow(o.timeout); err != nil {
    66  		return nil, err
    67  	}
    68  
    69  	return cs.Copy(), nil
    70  }