mongodb-util/src/qfw/mongodb/mongodbSim.go

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)
}