/mgo/socket.go
Go | 529 lines | 411 code | 63 blank | 55 comment | 87 complexity | 26e0bf4daff057c68ff22b6eb126c019 MD5 | raw file
Possible License(s): BSD-2-Clause
- // mgo - MongoDB driver for Go
- //
- // Copyright (c) 2010-2012 - Gustavo Niemeyer <gustavo@niemeyer.net>
- //
- // All rights reserved.
- //
- // Redistribution and use in source and binary forms, with or without
- // modification, are permitted provided that the following conditions are met:
- //
- // 1. Redistributions of source code must retain the above copyright notice, this
- // list of conditions and the following disclaimer.
- // 2. Redistributions in binary form must reproduce the above copyright notice,
- // this list of conditions and the following disclaimer in the documentation
- // and/or other materials provided with the distribution.
- //
- // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
- // ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
- // WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
- // ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
- // (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
- // LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
- // ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
- // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
- // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
- package mgo
- import (
- "errors"
- "unisearch/mgo/bson"
- "net"
- "sync"
- )
- type replyFunc func(err error, reply *replyOp, docNum int, docData []byte)
- type mongoSocket struct {
- sync.Mutex
- server *mongoServer // nil when cached
- conn *net.TCPConn
- addr string // For debugging only.
- nextRequestId uint32
- replyFuncs map[uint32]replyFunc
- references int
- auth []authInfo
- logout []authInfo
- cachedNonce string
- gotNonce sync.Cond
- dead error
- }
- type queryOp struct {
- collection string
- query interface{}
- skip int32
- limit int32
- selector interface{}
- flags uint32
- replyFunc replyFunc
- }
- type getMoreOp struct {
- collection string
- limit int32
- cursorId int64
- replyFunc replyFunc
- }
- type replyOp struct {
- flags uint32
- cursorId int64
- firstDoc int32
- replyDocs int32
- }
- type insertOp struct {
- collection string // "database.collection"
- documents []interface{} // One or more documents to insert
- }
- type updateOp struct {
- collection string // "database.collection"
- selector interface{}
- update interface{}
- flags uint32
- }
- type deleteOp struct {
- collection string // "database.collection"
- selector interface{}
- flags uint32
- }
- type requestInfo struct {
- bufferPos int
- replyFunc replyFunc
- }
- func newSocket(server *mongoServer, conn *net.TCPConn) *mongoSocket {
- socket := &mongoSocket{conn: conn, addr: server.Addr}
- socket.gotNonce.L = &socket.Mutex
- socket.replyFuncs = make(map[uint32]replyFunc)
- socket.server = server
- if err := socket.InitialAcquire(); err != nil {
- panic("newSocket: InitialAcquire returned error: " + err.Error())
- }
- stats.socketsAlive(+1)
- debugf("Socket %p to %s: initialized", socket, socket.addr)
- socket.resetNonce()
- go socket.readLoop()
- return socket
- }
- // InitialAcquire obtains the first reference to the socket, either
- // right after the connection is made or once a recycled socket is
- // being put back in use.
- func (socket *mongoSocket) InitialAcquire() error {
- socket.Lock()
- if socket.references > 0 {
- panic("Socket acquired out of cache with references")
- }
- if socket.dead != nil {
- socket.Unlock()
- return socket.dead
- }
- socket.references++
- stats.socketsInUse(+1)
- stats.socketRefs(+1)
- socket.Unlock()
- return nil
- }
- // Acquire obtains an additional reference to the socket.
- // The socket will only be recycled when it's released as many
- // times as it's been acquired.
- func (socket *mongoSocket) Acquire() (isMaster bool) {
- socket.Lock()
- if socket.references == 0 {
- panic("Socket got non-initial acquire with references == 0")
- }
- socket.references++
- stats.socketRefs(+1)
- // We'll track references to dead sockets as well.
- // Caller is still supposed to release the socket.
- if socket.dead == nil {
- isMaster = socket.server.IsMaster()
- }
- socket.Unlock()
- return isMaster
- }
- // Release decrements a socket reference. The socket will be
- // recycled once its released as many times as it's been acquired.
- func (socket *mongoSocket) Release() {
- socket.Lock()
- if socket.references == 0 {
- panic("socket.Release() with references == 0")
- }
- socket.references--
- stats.socketRefs(-1)
- if socket.references == 0 {
- stats.socketsInUse(-1)
- server := socket.server
- socket.Unlock()
- socket.LogoutAll()
- // If the socket is dead server is nil.
- if server != nil {
- server.RecycleSocket(socket)
- }
- } else {
- socket.Unlock()
- }
- }
- // Close terminates the socket use.
- func (socket *mongoSocket) Close() {
- socket.kill(errors.New("Closed explicitly"), false)
- }
- func (socket *mongoSocket) kill(err error, abend bool) {
- socket.Lock()
- if socket.dead != nil {
- debugf("Socket %p to %s: killed again: %s (previously: %s)", socket, socket.addr, err.Error(), socket.dead.Error())
- socket.Unlock()
- return
- }
- logf("Socket %p to %s: closing: %s (abend=%v)", socket, socket.addr, err.Error(), abend)
- socket.dead = err
- socket.conn.Close()
- stats.socketsAlive(-1)
- replyFuncs := socket.replyFuncs
- socket.replyFuncs = make(map[uint32]replyFunc)
- server := socket.server
- socket.server = nil
- socket.Unlock()
- for _, f := range replyFuncs {
- logf("Socket %p to %s: notifying replyFunc of closed socket: %s", socket, socket.addr, err.Error())
- f(err, nil, -1, nil)
- }
- if abend {
- server.AbendSocket(socket)
- }
- }
- func (socket *mongoSocket) SimpleQuery(op *queryOp) (data []byte, err error) {
- var mutex sync.Mutex
- var replyData []byte
- var replyErr error
- mutex.Lock()
- op.replyFunc = func(err error, reply *replyOp, docNum int, docData []byte) {
- replyData = docData
- replyErr = err
- mutex.Unlock()
- }
- err = socket.Query(op)
- if err != nil {
- return nil, err
- }
- mutex.Lock() // Wait.
- if replyErr != nil {
- return nil, replyErr
- }
- return replyData, nil
- }
- func (socket *mongoSocket) Query(ops ...interface{}) (err error) {
- if lops := socket.flushLogout(); len(lops) > 0 {
- ops = append(lops, ops...)
- }
- buf := make([]byte, 0, 256)
- // Serialize operations synchronously to avoid interrupting
- // other goroutines while we can't really be sending data.
- // Also, record id positions so that we can compute request
- // ids at once later with the lock already held.
- requests := make([]requestInfo, len(ops))
- requestCount := 0
- for _, op := range ops {
- debugf("Socket %p to %s: serializing op: %#v", socket, socket.addr, op)
- start := len(buf)
- var replyFunc replyFunc
- switch op := op.(type) {
- case *updateOp:
- buf = addHeader(buf, 2001)
- buf = addInt32(buf, 0) // Reserved
- buf = addCString(buf, op.collection)
- buf = addInt32(buf, int32(op.flags))
- debugf("Socket %p to %s: serializing selector document: %#v", socket, socket.addr, op.selector)
- buf, err = addBSON(buf, op.selector)
- if err != nil {
- return err
- }
- debugf("Socket %p to %s: serializing update document: %#v", socket, socket.addr, op.update)
- buf, err = addBSON(buf, op.update)
- if err != nil {
- return err
- }
- case *insertOp:
- buf = addHeader(buf, 2002)
- buf = addInt32(buf, 0) // Reserved
- buf = addCString(buf, op.collection)
- for _, doc := range op.documents {
- debugf("Socket %p to %s: serializing document for insertion: %#v", socket, socket.addr, doc)
- buf, err = addBSON(buf, doc)
- if err != nil {
- return err
- }
- }
- case *queryOp:
- buf = addHeader(buf, 2004)
- buf = addInt32(buf, int32(op.flags))
- buf = addCString(buf, op.collection)
- buf = addInt32(buf, op.skip)
- buf = addInt32(buf, op.limit)
- buf, err = addBSON(buf, op.query)
- if err != nil {
- return err
- }
- if op.selector != nil {
- buf, err = addBSON(buf, op.selector)
- if err != nil {
- return err
- }
- }
- replyFunc = op.replyFunc
- case *getMoreOp:
- buf = addHeader(buf, 2005)
- buf = addInt32(buf, 0) // Reserved
- buf = addCString(buf, op.collection)
- buf = addInt32(buf, op.limit)
- buf = addInt64(buf, op.cursorId)
- replyFunc = op.replyFunc
- case *deleteOp:
- buf = addHeader(buf, 2006)
- buf = addInt32(buf, 0) // Reserved
- buf = addCString(buf, op.collection)
- buf = addInt32(buf, int32(op.flags))
- debugf("Socket %p to %s: serializing selector document: %#v", socket, socket.addr, op.selector)
- buf, err = addBSON(buf, op.selector)
- if err != nil {
- return err
- }
- default:
- panic("Internal error: unknown operation type")
- }
- setInt32(buf, start, int32(len(buf)-start))
- if replyFunc != nil {
- request := &requests[requestCount]
- request.replyFunc = replyFunc
- request.bufferPos = start
- requestCount++
- }
- }
- // Buffer is ready for the pipe. Lock, allocate ids, and enqueue.
- socket.Lock()
- if socket.dead != nil {
- socket.Unlock()
- debug("Socket %p to %s: failing query, already closed: %s", socket, socket.addr, socket.dead.Error())
- // XXX This seems necessary in case the session is closed concurrently
- // with a query being performed, but it's not yet tested:
- for i := 0; i != requestCount; i++ {
- request := &requests[i]
- if request.replyFunc != nil {
- request.replyFunc(socket.dead, nil, -1, nil)
- }
- }
- return socket.dead
- }
- // Reserve id 0 for requests which should have no responses.
- requestId := socket.nextRequestId + 1
- if requestId == 0 {
- requestId++
- }
- socket.nextRequestId = requestId + uint32(requestCount)
- for i := 0; i != requestCount; i++ {
- request := &requests[i]
- setInt32(buf, request.bufferPos+4, int32(requestId))
- socket.replyFuncs[requestId] = request.replyFunc
- requestId++
- }
- debugf("Socket %p to %s: sending %d op(s) (%d bytes)", socket, socket.addr, len(ops), len(buf))
- stats.sentOps(len(ops))
- _, err = socket.conn.Write(buf)
- socket.Unlock()
- return err
- }
- func fill(r *net.TCPConn, b []byte) error {
- l := len(b)
- n, err := r.Read(b)
- for n != l && err == nil {
- var ni int
- ni, err = r.Read(b[n:])
- n += ni
- }
- return err
- }
- // Estimated minimum cost per socket: 1 goroutine + memory for the largest
- // document ever seen.
- func (socket *mongoSocket) readLoop() {
- p := make([]byte, 36) // 16 from header + 20 from OP_REPLY fixed fields
- s := make([]byte, 4)
- conn := socket.conn // No locking, conn never changes.
- for {
- // XXX Handle timeouts, , etc
- err := fill(conn, p)
- if err != nil {
- socket.kill(err, true)
- return
- }
- totalLen := getInt32(p, 0)
- responseTo := getInt32(p, 8)
- opCode := getInt32(p, 12)
- // Don't use socket.server.Addr here. socket is not
- // locked and socket.server may go away.
- debugf("Socket %p to %s: got reply (%d bytes)", socket, socket.addr, totalLen)
- _ = totalLen
- if opCode != 1 {
- socket.kill(errors.New("opcode != 1, corrupted data?"), true)
- return
- }
- reply := replyOp{
- flags: uint32(getInt32(p, 16)),
- cursorId: getInt64(p, 20),
- firstDoc: getInt32(p, 28),
- replyDocs: getInt32(p, 32),
- }
- stats.receivedOps(+1)
- stats.receivedDocs(int(reply.replyDocs))
- socket.Lock()
- replyFunc, replyFuncFound := socket.replyFuncs[uint32(responseTo)]
- socket.Unlock()
- if replyFunc != nil && reply.replyDocs == 0 {
- replyFunc(nil, &reply, -1, nil)
- } else {
- for i := 0; i != int(reply.replyDocs); i++ {
- err := fill(conn, s)
- if err != nil {
- socket.kill(err, true)
- return
- }
- b := make([]byte, int(getInt32(s, 0)))
- // copy(b, s) in an efficient way.
- b[0] = s[0]
- b[1] = s[1]
- b[2] = s[2]
- b[3] = s[3]
- err = fill(conn, b[4:])
- if err != nil {
- socket.kill(err, true)
- return
- }
- if globalDebug && globalLogger != nil {
- m := bson.M{}
- if err := bson.Unmarshal(b, m); err == nil {
- debugf("Socket %p to %s: received document: %#v", socket, socket.addr, m)
- }
- }
- if replyFunc != nil {
- replyFunc(nil, &reply, i, b)
- }
- // XXX Do bound checking against totalLen.
- }
- }
- // Only remove replyFunc after iteration, so that kill() will see it.
- socket.Lock()
- if replyFuncFound {
- delete(socket.replyFuncs, uint32(responseTo))
- }
- socket.Unlock()
- // XXX Do bound checking against totalLen.
- }
- }
- var emptyHeader = []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}
- func addHeader(b []byte, opcode int) []byte {
- i := len(b)
- b = append(b, emptyHeader...)
- // Enough for current opcodes.
- b[i+12] = byte(opcode)
- b[i+13] = byte(opcode >> 8)
- return b
- }
- func addInt32(b []byte, i int32) []byte {
- return append(b, byte(i), byte(i>>8), byte(i>>16), byte(i>>24))
- }
- func addInt64(b []byte, i int64) []byte {
- return append(b, byte(i), byte(i>>8), byte(i>>16), byte(i>>24),
- byte(i>>32), byte(i>>40), byte(i>>48), byte(i>>56))
- }
- func addCString(b []byte, s string) []byte {
- b = append(b, []byte(s)...)
- b = append(b, 0)
- return b
- }
- func addBSON(b []byte, doc interface{}) ([]byte, error) {
- if doc == nil {
- return append(b, 5, 0, 0, 0, 0), nil
- }
- data, err := bson.Marshal(doc)
- if err != nil {
- return b, err
- }
- return append(b, data...), nil
- }
- func setInt32(b []byte, pos int, i int32) {
- b[pos] = byte(i)
- b[pos+1] = byte(i >> 8)
- b[pos+2] = byte(i >> 16)
- b[pos+3] = byte(i >> 24)
- }
- func getInt32(b []byte, pos int) int32 {
- return (int32(b[pos+0])) |
- (int32(b[pos+1]) << 8) |
- (int32(b[pos+2]) << 16) |
- (int32(b[pos+3]) << 24)
- }
- func getInt64(b []byte, pos int) int64 {
- return (int64(b[pos+0])) |
- (int64(b[pos+1]) << 8) |
- (int64(b[pos+2]) << 16) |
- (int64(b[pos+3]) << 24) |
- (int64(b[pos+4]) << 32) |
- (int64(b[pos+5]) << 40) |
- (int64(b[pos+6]) << 48) |
- (int64(b[pos+7]) << 56)
- }