PageRenderTime 46ms CodeModel.GetById 20ms RepoModel.GetById 0ms app.codeStats 0ms

/mgo/socket.go

https://bitbucket.org/zuko_uno/unisearch
Go | 529 lines | 411 code | 63 blank | 55 comment | 87 complexity | 26e0bf4daff057c68ff22b6eb126c019 MD5 | raw file
Possible License(s): BSD-2-Clause
  1. // mgo - MongoDB driver for Go
  2. //
  3. // Copyright (c) 2010-2012 - Gustavo Niemeyer <gustavo@niemeyer.net>
  4. //
  5. // All rights reserved.
  6. //
  7. // Redistribution and use in source and binary forms, with or without
  8. // modification, are permitted provided that the following conditions are met:
  9. //
  10. // 1. Redistributions of source code must retain the above copyright notice, this
  11. // list of conditions and the following disclaimer.
  12. // 2. Redistributions in binary form must reproduce the above copyright notice,
  13. // this list of conditions and the following disclaimer in the documentation
  14. // and/or other materials provided with the distribution.
  15. //
  16. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
  17. // ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  18. // WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  19. // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
  20. // ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
  21. // (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
  22. // LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
  23. // ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  24. // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  25. // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  26. package mgo
  27. import (
  28. "errors"
  29. "unisearch/mgo/bson"
  30. "net"
  31. "sync"
  32. )
  33. type replyFunc func(err error, reply *replyOp, docNum int, docData []byte)
  34. type mongoSocket struct {
  35. sync.Mutex
  36. server *mongoServer // nil when cached
  37. conn *net.TCPConn
  38. addr string // For debugging only.
  39. nextRequestId uint32
  40. replyFuncs map[uint32]replyFunc
  41. references int
  42. auth []authInfo
  43. logout []authInfo
  44. cachedNonce string
  45. gotNonce sync.Cond
  46. dead error
  47. }
  48. type queryOp struct {
  49. collection string
  50. query interface{}
  51. skip int32
  52. limit int32
  53. selector interface{}
  54. flags uint32
  55. replyFunc replyFunc
  56. }
  57. type getMoreOp struct {
  58. collection string
  59. limit int32
  60. cursorId int64
  61. replyFunc replyFunc
  62. }
  63. type replyOp struct {
  64. flags uint32
  65. cursorId int64
  66. firstDoc int32
  67. replyDocs int32
  68. }
  69. type insertOp struct {
  70. collection string // "database.collection"
  71. documents []interface{} // One or more documents to insert
  72. }
  73. type updateOp struct {
  74. collection string // "database.collection"
  75. selector interface{}
  76. update interface{}
  77. flags uint32
  78. }
  79. type deleteOp struct {
  80. collection string // "database.collection"
  81. selector interface{}
  82. flags uint32
  83. }
  84. type requestInfo struct {
  85. bufferPos int
  86. replyFunc replyFunc
  87. }
  88. func newSocket(server *mongoServer, conn *net.TCPConn) *mongoSocket {
  89. socket := &mongoSocket{conn: conn, addr: server.Addr}
  90. socket.gotNonce.L = &socket.Mutex
  91. socket.replyFuncs = make(map[uint32]replyFunc)
  92. socket.server = server
  93. if err := socket.InitialAcquire(); err != nil {
  94. panic("newSocket: InitialAcquire returned error: " + err.Error())
  95. }
  96. stats.socketsAlive(+1)
  97. debugf("Socket %p to %s: initialized", socket, socket.addr)
  98. socket.resetNonce()
  99. go socket.readLoop()
  100. return socket
  101. }
  102. // InitialAcquire obtains the first reference to the socket, either
  103. // right after the connection is made or once a recycled socket is
  104. // being put back in use.
  105. func (socket *mongoSocket) InitialAcquire() error {
  106. socket.Lock()
  107. if socket.references > 0 {
  108. panic("Socket acquired out of cache with references")
  109. }
  110. if socket.dead != nil {
  111. socket.Unlock()
  112. return socket.dead
  113. }
  114. socket.references++
  115. stats.socketsInUse(+1)
  116. stats.socketRefs(+1)
  117. socket.Unlock()
  118. return nil
  119. }
  120. // Acquire obtains an additional reference to the socket.
  121. // The socket will only be recycled when it's released as many
  122. // times as it's been acquired.
  123. func (socket *mongoSocket) Acquire() (isMaster bool) {
  124. socket.Lock()
  125. if socket.references == 0 {
  126. panic("Socket got non-initial acquire with references == 0")
  127. }
  128. socket.references++
  129. stats.socketRefs(+1)
  130. // We'll track references to dead sockets as well.
  131. // Caller is still supposed to release the socket.
  132. if socket.dead == nil {
  133. isMaster = socket.server.IsMaster()
  134. }
  135. socket.Unlock()
  136. return isMaster
  137. }
  138. // Release decrements a socket reference. The socket will be
  139. // recycled once its released as many times as it's been acquired.
  140. func (socket *mongoSocket) Release() {
  141. socket.Lock()
  142. if socket.references == 0 {
  143. panic("socket.Release() with references == 0")
  144. }
  145. socket.references--
  146. stats.socketRefs(-1)
  147. if socket.references == 0 {
  148. stats.socketsInUse(-1)
  149. server := socket.server
  150. socket.Unlock()
  151. socket.LogoutAll()
  152. // If the socket is dead server is nil.
  153. if server != nil {
  154. server.RecycleSocket(socket)
  155. }
  156. } else {
  157. socket.Unlock()
  158. }
  159. }
  160. // Close terminates the socket use.
  161. func (socket *mongoSocket) Close() {
  162. socket.kill(errors.New("Closed explicitly"), false)
  163. }
  164. func (socket *mongoSocket) kill(err error, abend bool) {
  165. socket.Lock()
  166. if socket.dead != nil {
  167. debugf("Socket %p to %s: killed again: %s (previously: %s)", socket, socket.addr, err.Error(), socket.dead.Error())
  168. socket.Unlock()
  169. return
  170. }
  171. logf("Socket %p to %s: closing: %s (abend=%v)", socket, socket.addr, err.Error(), abend)
  172. socket.dead = err
  173. socket.conn.Close()
  174. stats.socketsAlive(-1)
  175. replyFuncs := socket.replyFuncs
  176. socket.replyFuncs = make(map[uint32]replyFunc)
  177. server := socket.server
  178. socket.server = nil
  179. socket.Unlock()
  180. for _, f := range replyFuncs {
  181. logf("Socket %p to %s: notifying replyFunc of closed socket: %s", socket, socket.addr, err.Error())
  182. f(err, nil, -1, nil)
  183. }
  184. if abend {
  185. server.AbendSocket(socket)
  186. }
  187. }
  188. func (socket *mongoSocket) SimpleQuery(op *queryOp) (data []byte, err error) {
  189. var mutex sync.Mutex
  190. var replyData []byte
  191. var replyErr error
  192. mutex.Lock()
  193. op.replyFunc = func(err error, reply *replyOp, docNum int, docData []byte) {
  194. replyData = docData
  195. replyErr = err
  196. mutex.Unlock()
  197. }
  198. err = socket.Query(op)
  199. if err != nil {
  200. return nil, err
  201. }
  202. mutex.Lock() // Wait.
  203. if replyErr != nil {
  204. return nil, replyErr
  205. }
  206. return replyData, nil
  207. }
  208. func (socket *mongoSocket) Query(ops ...interface{}) (err error) {
  209. if lops := socket.flushLogout(); len(lops) > 0 {
  210. ops = append(lops, ops...)
  211. }
  212. buf := make([]byte, 0, 256)
  213. // Serialize operations synchronously to avoid interrupting
  214. // other goroutines while we can't really be sending data.
  215. // Also, record id positions so that we can compute request
  216. // ids at once later with the lock already held.
  217. requests := make([]requestInfo, len(ops))
  218. requestCount := 0
  219. for _, op := range ops {
  220. debugf("Socket %p to %s: serializing op: %#v", socket, socket.addr, op)
  221. start := len(buf)
  222. var replyFunc replyFunc
  223. switch op := op.(type) {
  224. case *updateOp:
  225. buf = addHeader(buf, 2001)
  226. buf = addInt32(buf, 0) // Reserved
  227. buf = addCString(buf, op.collection)
  228. buf = addInt32(buf, int32(op.flags))
  229. debugf("Socket %p to %s: serializing selector document: %#v", socket, socket.addr, op.selector)
  230. buf, err = addBSON(buf, op.selector)
  231. if err != nil {
  232. return err
  233. }
  234. debugf("Socket %p to %s: serializing update document: %#v", socket, socket.addr, op.update)
  235. buf, err = addBSON(buf, op.update)
  236. if err != nil {
  237. return err
  238. }
  239. case *insertOp:
  240. buf = addHeader(buf, 2002)
  241. buf = addInt32(buf, 0) // Reserved
  242. buf = addCString(buf, op.collection)
  243. for _, doc := range op.documents {
  244. debugf("Socket %p to %s: serializing document for insertion: %#v", socket, socket.addr, doc)
  245. buf, err = addBSON(buf, doc)
  246. if err != nil {
  247. return err
  248. }
  249. }
  250. case *queryOp:
  251. buf = addHeader(buf, 2004)
  252. buf = addInt32(buf, int32(op.flags))
  253. buf = addCString(buf, op.collection)
  254. buf = addInt32(buf, op.skip)
  255. buf = addInt32(buf, op.limit)
  256. buf, err = addBSON(buf, op.query)
  257. if err != nil {
  258. return err
  259. }
  260. if op.selector != nil {
  261. buf, err = addBSON(buf, op.selector)
  262. if err != nil {
  263. return err
  264. }
  265. }
  266. replyFunc = op.replyFunc
  267. case *getMoreOp:
  268. buf = addHeader(buf, 2005)
  269. buf = addInt32(buf, 0) // Reserved
  270. buf = addCString(buf, op.collection)
  271. buf = addInt32(buf, op.limit)
  272. buf = addInt64(buf, op.cursorId)
  273. replyFunc = op.replyFunc
  274. case *deleteOp:
  275. buf = addHeader(buf, 2006)
  276. buf = addInt32(buf, 0) // Reserved
  277. buf = addCString(buf, op.collection)
  278. buf = addInt32(buf, int32(op.flags))
  279. debugf("Socket %p to %s: serializing selector document: %#v", socket, socket.addr, op.selector)
  280. buf, err = addBSON(buf, op.selector)
  281. if err != nil {
  282. return err
  283. }
  284. default:
  285. panic("Internal error: unknown operation type")
  286. }
  287. setInt32(buf, start, int32(len(buf)-start))
  288. if replyFunc != nil {
  289. request := &requests[requestCount]
  290. request.replyFunc = replyFunc
  291. request.bufferPos = start
  292. requestCount++
  293. }
  294. }
  295. // Buffer is ready for the pipe. Lock, allocate ids, and enqueue.
  296. socket.Lock()
  297. if socket.dead != nil {
  298. socket.Unlock()
  299. debug("Socket %p to %s: failing query, already closed: %s", socket, socket.addr, socket.dead.Error())
  300. // XXX This seems necessary in case the session is closed concurrently
  301. // with a query being performed, but it's not yet tested:
  302. for i := 0; i != requestCount; i++ {
  303. request := &requests[i]
  304. if request.replyFunc != nil {
  305. request.replyFunc(socket.dead, nil, -1, nil)
  306. }
  307. }
  308. return socket.dead
  309. }
  310. // Reserve id 0 for requests which should have no responses.
  311. requestId := socket.nextRequestId + 1
  312. if requestId == 0 {
  313. requestId++
  314. }
  315. socket.nextRequestId = requestId + uint32(requestCount)
  316. for i := 0; i != requestCount; i++ {
  317. request := &requests[i]
  318. setInt32(buf, request.bufferPos+4, int32(requestId))
  319. socket.replyFuncs[requestId] = request.replyFunc
  320. requestId++
  321. }
  322. debugf("Socket %p to %s: sending %d op(s) (%d bytes)", socket, socket.addr, len(ops), len(buf))
  323. stats.sentOps(len(ops))
  324. _, err = socket.conn.Write(buf)
  325. socket.Unlock()
  326. return err
  327. }
  328. func fill(r *net.TCPConn, b []byte) error {
  329. l := len(b)
  330. n, err := r.Read(b)
  331. for n != l && err == nil {
  332. var ni int
  333. ni, err = r.Read(b[n:])
  334. n += ni
  335. }
  336. return err
  337. }
  338. // Estimated minimum cost per socket: 1 goroutine + memory for the largest
  339. // document ever seen.
  340. func (socket *mongoSocket) readLoop() {
  341. p := make([]byte, 36) // 16 from header + 20 from OP_REPLY fixed fields
  342. s := make([]byte, 4)
  343. conn := socket.conn // No locking, conn never changes.
  344. for {
  345. // XXX Handle timeouts, , etc
  346. err := fill(conn, p)
  347. if err != nil {
  348. socket.kill(err, true)
  349. return
  350. }
  351. totalLen := getInt32(p, 0)
  352. responseTo := getInt32(p, 8)
  353. opCode := getInt32(p, 12)
  354. // Don't use socket.server.Addr here. socket is not
  355. // locked and socket.server may go away.
  356. debugf("Socket %p to %s: got reply (%d bytes)", socket, socket.addr, totalLen)
  357. _ = totalLen
  358. if opCode != 1 {
  359. socket.kill(errors.New("opcode != 1, corrupted data?"), true)
  360. return
  361. }
  362. reply := replyOp{
  363. flags: uint32(getInt32(p, 16)),
  364. cursorId: getInt64(p, 20),
  365. firstDoc: getInt32(p, 28),
  366. replyDocs: getInt32(p, 32),
  367. }
  368. stats.receivedOps(+1)
  369. stats.receivedDocs(int(reply.replyDocs))
  370. socket.Lock()
  371. replyFunc, replyFuncFound := socket.replyFuncs[uint32(responseTo)]
  372. socket.Unlock()
  373. if replyFunc != nil && reply.replyDocs == 0 {
  374. replyFunc(nil, &reply, -1, nil)
  375. } else {
  376. for i := 0; i != int(reply.replyDocs); i++ {
  377. err := fill(conn, s)
  378. if err != nil {
  379. socket.kill(err, true)
  380. return
  381. }
  382. b := make([]byte, int(getInt32(s, 0)))
  383. // copy(b, s) in an efficient way.
  384. b[0] = s[0]
  385. b[1] = s[1]
  386. b[2] = s[2]
  387. b[3] = s[3]
  388. err = fill(conn, b[4:])
  389. if err != nil {
  390. socket.kill(err, true)
  391. return
  392. }
  393. if globalDebug && globalLogger != nil {
  394. m := bson.M{}
  395. if err := bson.Unmarshal(b, m); err == nil {
  396. debugf("Socket %p to %s: received document: %#v", socket, socket.addr, m)
  397. }
  398. }
  399. if replyFunc != nil {
  400. replyFunc(nil, &reply, i, b)
  401. }
  402. // XXX Do bound checking against totalLen.
  403. }
  404. }
  405. // Only remove replyFunc after iteration, so that kill() will see it.
  406. socket.Lock()
  407. if replyFuncFound {
  408. delete(socket.replyFuncs, uint32(responseTo))
  409. }
  410. socket.Unlock()
  411. // XXX Do bound checking against totalLen.
  412. }
  413. }
  414. var emptyHeader = []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}
  415. func addHeader(b []byte, opcode int) []byte {
  416. i := len(b)
  417. b = append(b, emptyHeader...)
  418. // Enough for current opcodes.
  419. b[i+12] = byte(opcode)
  420. b[i+13] = byte(opcode >> 8)
  421. return b
  422. }
  423. func addInt32(b []byte, i int32) []byte {
  424. return append(b, byte(i), byte(i>>8), byte(i>>16), byte(i>>24))
  425. }
  426. func addInt64(b []byte, i int64) []byte {
  427. return append(b, byte(i), byte(i>>8), byte(i>>16), byte(i>>24),
  428. byte(i>>32), byte(i>>40), byte(i>>48), byte(i>>56))
  429. }
  430. func addCString(b []byte, s string) []byte {
  431. b = append(b, []byte(s)...)
  432. b = append(b, 0)
  433. return b
  434. }
  435. func addBSON(b []byte, doc interface{}) ([]byte, error) {
  436. if doc == nil {
  437. return append(b, 5, 0, 0, 0, 0), nil
  438. }
  439. data, err := bson.Marshal(doc)
  440. if err != nil {
  441. return b, err
  442. }
  443. return append(b, data...), nil
  444. }
  445. func setInt32(b []byte, pos int, i int32) {
  446. b[pos] = byte(i)
  447. b[pos+1] = byte(i >> 8)
  448. b[pos+2] = byte(i >> 16)
  449. b[pos+3] = byte(i >> 24)
  450. }
  451. func getInt32(b []byte, pos int) int32 {
  452. return (int32(b[pos+0])) |
  453. (int32(b[pos+1]) << 8) |
  454. (int32(b[pos+2]) << 16) |
  455. (int32(b[pos+3]) << 24)
  456. }
  457. func getInt64(b []byte, pos int) int64 {
  458. return (int64(b[pos+0])) |
  459. (int64(b[pos+1]) << 8) |
  460. (int64(b[pos+2]) << 16) |
  461. (int64(b[pos+3]) << 24) |
  462. (int64(b[pos+4]) << 32) |
  463. (int64(b[pos+5]) << 40) |
  464. (int64(b[pos+6]) << 48) |
  465. (int64(b[pos+7]) << 56)
  466. }