package mongodb import ( "errors" "log" "strings" "time" mgo "gopkg.in/mgo.v2" . "gopkg.in/mgo.v2/bson" ) type MongodbSim struct { pool chan *mgo.Session globalS *mgo.Session idlePool chan bool MongodbAddr string Size int Half int //初始化一半的连接 DbName string ReplSet string } func NewMongodbPool(size int, addr, dbName string) *MongodbSim { return NewMongodbReplPool(size, addr, dbName, "") } func NewMongodbReplPool(size int, addr, dbName, replset string) (m *MongodbSim) { m = &MongodbSim{ MongodbAddr: addr, DbName: dbName, Size: size, ReplSet: replset, } m.InitPool() return } //全部关闭,一般不用 func (m *MongodbSim) Destory() { L1: for { select { case session := <-m.pool: if session != nil { session.Close() session = nil } case <-time.After(time.Millisecond * 50): break L1 } } } func (m *MongodbSim) InitPool() { defer Catch() ms := strings.Split(m.MongodbAddr, ",") if m.ReplSet == "" && len(ms) > 1 { m.ReplSet = "qfws" } m.Half = m.Size/2 + 1 m.pool = make(chan *mgo.Session, m.Size) m.idlePool = make(chan bool, m.Size) for i := 0; i < m.Size; i++ { if i < m.Half { sess, err := m.createConn() if err == nil && sess.Ping() == nil { m.pool <- sess } } else { m.idlePool <- true } } go m.GC() } //清理多余连接 func (m *MongodbSim) GC() { for { i, n := 0, 0 if len(m.pool) > m.Half { //开始清理多余连接,不考虑实时动态变化,保留5分之一连接 i = len(m.pool) - (m.Size/5 + 1) if i > 0 { L: for ; n < i; n++ { select { case sess := <-m.pool: sess.SetSocketTimeout(1) sess.SetSyncTimeout(1) sess.Close() sess = nil select { case m.idlePool <- true: default: } default: break L } } } if n > 0 { log.Println("Sim-pool,idlepool:", len(m.pool), len(m.idlePool), "预计清除,实际清除", i, n) } } time.Sleep(1 * time.Minute) } } func (m *MongodbSim) GetMgoConn(sec ...int) (session *mgo.Session) { defer Catch() select { case session = <-m.pool: err := session.Ping() if err != nil { session.Close() session, _ = m.createConn() } if session == nil { select { case m.idlePool <- true: default: } } default: select { case session = <-m.pool: err := session.Ping() if err != nil { session.Close() session, _ = m.createConn() } if session == nil { select { case m.idlePool <- true: default: } } case <-m.idlePool: //空闲 session, _ = m.createConn() case <-time.After(3 * time.Second): //超时 session, _ = m.createConn() } } timeout := TIMEOUT if session != nil && len(sec) == 1 && sec[0] > 4 { timeout = time.Duration(sec[0]) * time.Second session.SetSocketTimeout(timeout) session.SetSyncTimeout(timeout) } return } func (m *MongodbSim) DestoryMongoConn(session *mgo.Session) { select { case m.pool <- session: default: } } //取session链接 func (m *MongodbSim) createConn() (sess *mgo.Session, err error) { //增加对集群的支持 if m.globalS == nil { info := mgo.DialInfo{ Addrs: strings.Split(m.MongodbAddr, ","), Timeout: TIMEOUT, } info.SetMaxIdleTimeMS(120000) if ReplSet != "" { info.ReplicaSetName = m.ReplSet info.Direct = false } m.globalS, err = mgo.DialWithInfo(&info) if err != nil { log.Print("init mgo-global session", err) } } sess = m.globalS.Copy() return } //统计 func (m *MongodbSim) Count(c string, query interface{}) int { defer Catch() sess := m.GetMgoConn() //log.Println("count:", m.Size, m.MongodbAddr, m.DbName, sess, m.GetMgoConn(), m) var n int = 0 if sess != nil { defer m.DestoryMongoConn(sess) coll := sess.DB(m.DbName).C(c) var err error n, err = coll.Find(ObjToM(query)).Count() if nil != err { log.Println("CountError", err) } } return n } //统计 func (m *MongodbSim) CountByErr(c string, query interface{}) (int, error) { defer Catch() sess := m.GetMgoConn() //log.Println("count:", m.Size, m.MongodbAddr, m.DbName, sess, m.GetMgoConn(), m) var n int = 0 if sess != nil { defer m.DestoryMongoConn(sess) coll := sess.DB(m.DbName).C(c) var err error n, err = coll.Find(ObjToM(query)).Count() if nil != err { return 0, err } else { return n, nil } } return n, errors.New("no sess") } func (m *MongodbSim) Update(c string, query interface{}, set interface{}, upsert bool, multi bool) bool { defer Catch() sess := m.GetMgoConn() b := true if sess != nil { defer m.DestoryMongoConn(sess) coll := sess.DB(m.DbName).C(c) var err error if upsert { _, err = coll.Upsert(ObjToM(query), ObjToM(set)) } else { if multi { _, err = coll.UpdateAll(ObjToM(query), ObjToM(set)) } else { err = coll.Update(ObjToM(query), ObjToM(set)) } } if nil != err { log.Println("UpdateError", err) b = false } } return b } func (m *MongodbSim) UpdateById(c string, id interface{}, set interface{}) bool { defer Catch() sess := m.GetMgoConn() b := false if sess != nil { defer m.DestoryMongoConn(sess) coll := sess.DB(m.DbName).C(c) var q interface{} if sid, ok := id.(string); ok { q = M{"_id": StringTOBsonId(sid)} } else { q = M{"_id": id} } err := coll.Update(q, ObjToM(set)) if nil != err { log.Println("UpdateByIdError", err) b = false } else { b = true } } return b } //批量更新 func (m *MongodbSim) UpdateBulkAll(db, c string, doc ...[]map[string]interface{}) bool { defer Catch() sess := m.GetMgoConn() b := true if sess != nil { defer m.DestoryMongoConn(sess) coll := sess.DB(db).C(c) bulk := coll.Bulk() for _, v := range doc { if len(v) == 2 { bulk.Update(v[0], v[1]) } } _, err := bulk.Run() if nil != err { log.Println("BulkError", err) b = false } } else { b = false } return b } //新的批量修改 func (m *MongodbSim) NewUpdateBulk(c string, upsert bool, multi bool, doc ...[]map[string]interface{}) bool { defer Catch() sess := m.GetMgoConn() if sess == nil { return false } defer m.DestoryMongoConn(sess) coll := sess.DB(m.DbName).C(c) bulk := coll.Bulk() for _, v := range doc { if len(v) == 2 { if upsert { bulk.Upsert(v[0], v[1]) } else { if multi { bulk.UpdateAll(v[0], v[1]) } else { bulk.Update(v[0], v[1]) } } } } _, err := bulk.Run() if nil != err { log.Println("NewUpdateBulk", err) return false } return true } func (m *MongodbSim) UpdateBulk(c string, doc ...[]map[string]interface{}) bool { return m.UpdateBulkAll(m.DbName, c, doc...) } //批量更新 func (m *MongodbSim) UpSertBulk(c string, doc ...[]map[string]interface{}) bool { defer Catch() sess := m.GetMgoConn() b := true if sess != nil { defer m.DestoryMongoConn(sess) coll := sess.DB(m.DbName).C(c) bulk := coll.Bulk() for _, v := range doc { if len(v) == 2 { bulk.Upsert(v[0], v[1]) } } _, err := bulk.Run() if nil != err { log.Println("BulkUpsertError", err) b = false } } else { b = false } return b } //批量插入 func (m *MongodbSim) SaveBulk(c string, doc ...map[string]interface{}) bool { defer Catch() sess := m.GetMgoConn() b := true if sess != nil { defer m.DestoryMongoConn(sess) coll := sess.DB(m.DbName).C(c) bulk := coll.Bulk() for _, v := range doc { bulk.Insert(v) } _, err := bulk.Run() if nil != err { log.Println("BulkError", err) b = false } } else { b = false } return b } //批量插入 func (m *MongodbSim) SaveBulkInterface(c string, doc ...interface{}) bool { defer Catch() sess := m.GetMgoConn() b := true if sess != nil { defer m.DestoryMongoConn(sess) coll := sess.DB(m.DbName).C(c) bulk := coll.Bulk() for _, v := range doc { bulk.Insert(v) } _, err := bulk.Run() if nil != err { log.Println("BulkError", err) b = false } } else { b = false } return b } func (m *MongodbSim) Save(c string, doc interface{}) string { defer Catch() sess := m.GetMgoConn() if sess != nil { defer m.DestoryMongoConn(sess) coll := sess.DB(m.DbName).C(c) obj := ObjToM(doc) id := NewObjectId() (*obj)["_id"] = id err := coll.Insert(obj) if nil != err { log.Println("SaveError", err) return "" } return BsonIdToSId(id) } return "" } //原_id不变 func (m *MongodbSim) SaveByOriID(c string, doc interface{}) bool { defer Catch() sess := m.GetMgoConn() if sess != nil { defer m.DestoryMongoConn(sess) coll := sess.DB(m.DbName).C(c) obj := ObjToM(doc) err := coll.Insert(obj) if nil != err { log.Println("SaveError", err) return false } return true } return false } //查询单条对象 func (m *MongodbSim) FindOne(c string, query interface{}) (*map[string]interface{}, bool) { return m.FindOneByField(c, query, nil) } //查询单条对象 func (m *MongodbSim) FindOneByField(c string, query interface{}, fields interface{}) (*map[string]interface{}, bool) { defer Catch() res, ok := m.Find(c, query, nil, fields, true, -1, -1) if nil != res && len(*res) > 0 { return &((*res)[0]), ok } return nil, ok } //查询单条对象 func (m *MongodbSim) FindById(c string, query string, fields interface{}) (*map[string]interface{}, bool) { defer Catch() sess := m.GetMgoConn() var res map[string]interface{} b := false if sess != nil { defer m.DestoryMongoConn(sess) res = make(map[string]interface{}) coll := sess.DB(m.DbName).C(c) var err error err = coll.FindId(ObjectIdHex(query)).Select(ObjToOth(fields)).One(&res) if nil != err { log.Println("FindByIdError", err) } b = true } return &res, b } //底层查询方法 func (m *MongodbSim) Find(c string, query interface{}, order interface{}, fields interface{}, single bool, start int, limit int) (*[]map[string]interface{}, bool) { defer Catch() sess := m.GetMgoConn() var res []map[string]interface{} b := false if sess != nil { defer m.DestoryMongoConn(sess) res = make([]map[string]interface{}, 1) coll := sess.DB(m.DbName).C(c) var err error if single { err = coll.Find(ObjToM(query)).Select(ObjToOth(fields)).Sort(ObjToArr(order)...).One(&res[0]) } else if start > -1 { err = coll.Find(ObjToM(query)).Select(ObjToOth(fields)).Sort(ObjToArr(order)...).Skip(start).Limit(limit).All(&res) } else { err = coll.Find(ObjToM(query)).Select(ObjToOth(fields)).Sort(ObjToArr(order)...).All(&res) } if nil != err { //log.Println("FindError", err) } b = true } return &res, b } //删除对象 func (m *MongodbSim) Del(c string, query interface{}) bool { defer Catch() sess := m.GetMgoConn() b := false if sess != nil { defer m.DestoryMongoConn(sess) coll := sess.DB(m.DbName).C(c) _, err := coll.RemoveAll(ObjToM(query)) if nil != err { log.Println("DelError", err) b = false } else { b = true } } return b } func (m *MongodbSim) GetObjectId(str string) ObjectId { return ObjectIdHex(str) }