[
  {
    "path": ".travis.yml",
    "content": "language: go\nenv: GO_RUN_LONG_TEST=1\ngo:\n  - 1.5\ninstall:\n  - go get -t ./...\n"
  },
  {
    "path": "cmd/dvara/logger.go",
    "content": "package main\n\nimport \"log\"\n\n// stdLogger provides a logger backed by the standard library logger. This is a\n// placeholder until we can open source our logger.\ntype stdLogger struct{}\n\nfunc (l *stdLogger) Error(args ...interface{})                 { log.Print(args...) }\nfunc (l *stdLogger) Errorf(format string, args ...interface{}) { log.Printf(format, args...) }\nfunc (l *stdLogger) Warn(args ...interface{})                  { log.Print(args...) }\nfunc (l *stdLogger) Warnf(format string, args ...interface{})  { log.Printf(format, args...) }\nfunc (l *stdLogger) Info(args ...interface{})                  { log.Print(args...) }\nfunc (l *stdLogger) Infof(format string, args ...interface{})  { log.Printf(format, args...) }\nfunc (l *stdLogger) Debug(args ...interface{})                 { log.Print(args...) }\nfunc (l *stdLogger) Debugf(format string, args ...interface{}) { log.Printf(format, args...) }\n"
  },
  {
    "path": "cmd/dvara/main.go",
    "content": "package main\n\nimport (\n\t\"flag\"\n\t\"fmt\"\n\t\"os\"\n\t\"os/signal\"\n\t\"syscall\"\n\t\"time\"\n\n\t\"github.com/facebookgo/dvara\"\n\t\"github.com/facebookgo/inject\"\n\t\"github.com/facebookgo/startstop\"\n\t\"github.com/facebookgo/stats\"\n)\n\nfunc main() {\n\tif err := Main(); err != nil {\n\t\tfmt.Fprintln(os.Stderr, err)\n\t\tos.Exit(1)\n\t}\n}\n\nfunc Main() error {\n\tmessageTimeout := flag.Duration(\"message_timeout\", 2*time.Minute, \"timeout for one message to be proxied\")\n\tclientIdleTimeout := flag.Duration(\"client_idle_timeout\", 60*time.Minute, \"idle timeout for client connections\")\n\tserverIdleTimeout := flag.Duration(\"server_idle_timeout\", 1*time.Hour, \"idle timeout for  server connections\")\n\tserverClosePoolSize := flag.Uint(\"server_close_pool_size\", 100, \"number of goroutines that will handle closing server connections\")\n\tgetLastErrorTimeout := flag.Duration(\"get_last_error_timeout\", time.Minute, \"timeout for getLastError pinning\")\n\tmaxPerClientConnections := flag.Uint(\"max_per_client_connections\", 100, \"maximum number of connections per client\")\n\tmaxConnections := flag.Uint(\"max_connections\", 100, \"maximum number of connections per mongo\")\n\tportStart := flag.Int(\"port_start\", 6000, \"start of port range\")\n\tportEnd := flag.Int(\"port_end\", 6010, \"end of port range\")\n\taddrs := flag.String(\"addrs\", \"localhost:27017\", \"comma separated list of mongo addresses\")\n\n\tflag.Parse()\n\n\treplicaSet := dvara.ReplicaSet{\n\t\tAddrs:                   *addrs,\n\t\tPortStart:               *portStart,\n\t\tPortEnd:                 *portEnd,\n\t\tMessageTimeout:          *messageTimeout,\n\t\tClientIdleTimeout:       *clientIdleTimeout,\n\t\tServerIdleTimeout:       *serverIdleTimeout,\n\t\tServerClosePoolSize:     *serverClosePoolSize,\n\t\tGetLastErrorTimeout:     *getLastErrorTimeout,\n\t\tMaxConnections:          *maxConnections,\n\t\tMaxPerClientConnections: *maxPerClientConnections,\n\t}\n\n\tvar statsClient stats.HookClient\n\tvar log stdLogger\n\tvar graph inject.Graph\n\terr := graph.Provide(\n\t\t&inject.Object{Value: &log},\n\t\t&inject.Object{Value: &replicaSet},\n\t\t&inject.Object{Value: &statsClient},\n\t)\n\tif err != nil {\n\t\treturn err\n\t}\n\tif err := graph.Populate(); err != nil {\n\t\treturn err\n\t}\n\tobjects := graph.Objects()\n\n\tif err := startstop.Start(objects, &log); err != nil {\n\t\treturn err\n\t}\n\tdefer startstop.Stop(objects, &log)\n\n\tch := make(chan os.Signal, 2)\n\tsignal.Notify(ch, syscall.SIGTERM, syscall.SIGINT)\n\t<-ch\n\tsignal.Stop(ch)\n\treturn nil\n}\n"
  },
  {
    "path": "common_test.go",
    "content": "package dvara\n\nimport (\n\t\"os\"\n\t\"testing\"\n\t\"time\"\n\n\t\"gopkg.in/mgo.v2\"\n\n\t\"github.com/facebookgo/ensure\"\n\t\"github.com/facebookgo/inject\"\n\t\"github.com/facebookgo/mgotest\"\n\t\"github.com/facebookgo/startstop\"\n\t\"github.com/facebookgo/stats\"\n)\n\nvar (\n\tdisableSlowTests = os.Getenv(\"GO_RUN_LONG_TEST\") == \"\"\n\tveryVerbose      = os.Getenv(\"VERY_VERBOSE\") == \"1\"\n)\n\ntype tLogger struct {\n\tTB testing.TB\n}\n\nfunc (l *tLogger) Error(args ...interface{}) {\n\tif veryVerbose {\n\t\tl.TB.Log(args...)\n\t}\n}\n\nfunc (l *tLogger) Errorf(format string, args ...interface{}) {\n\tif veryVerbose {\n\t\tl.TB.Logf(format, args...)\n\t}\n}\n\nfunc (l *tLogger) Warn(args ...interface{}) {\n\tif veryVerbose {\n\t\tl.TB.Log(args...)\n\t}\n}\n\nfunc (l *tLogger) Warnf(format string, args ...interface{}) {\n\tif veryVerbose {\n\t\tl.TB.Logf(format, args...)\n\t}\n}\n\nfunc (l *tLogger) Info(args ...interface{}) {\n\tif veryVerbose {\n\t\tl.TB.Log(args...)\n\t}\n}\n\nfunc (l *tLogger) Infof(format string, args ...interface{}) {\n\tif veryVerbose {\n\t\tl.TB.Logf(format, args...)\n\t}\n}\n\nfunc (l *tLogger) Debug(args ...interface{}) {\n\tif veryVerbose {\n\t\tl.TB.Log(args...)\n\t}\n}\n\nfunc (l *tLogger) Debugf(format string, args ...interface{}) {\n\tif veryVerbose {\n\t\tl.TB.Logf(format, args...)\n\t}\n}\n\ntype stopper interface {\n\tStop()\n}\n\ntype Harness struct {\n\tT          testing.TB\n\tStopper    stopper // This is either mgotest.Server or mgotest.ReplicaSet\n\tReplicaSet *ReplicaSet\n\tGraph      *inject.Graph\n\tLog        *tLogger\n}\n\nfunc newHarnessInternal(url string, s stopper, t testing.TB) *Harness {\n\treplicaSet := ReplicaSet{\n\t\tAddrs:                   url,\n\t\tPortStart:               2000,\n\t\tPortEnd:                 3000,\n\t\tMaxConnections:          5,\n\t\tMinIdleConnections:      5,\n\t\tServerIdleTimeout:       5 * time.Minute,\n\t\tServerClosePoolSize:     5,\n\t\tClientIdleTimeout:       5 * time.Minute,\n\t\tMaxPerClientConnections: 250,\n\t\tGetLastErrorTimeout:     5 * time.Minute,\n\t\tMessageTimeout:          time.Minute,\n\t}\n\tlog := tLogger{TB: t}\n\tvar graph inject.Graph\n\terr := graph.Provide(\n\t\t&inject.Object{Value: &log},\n\t\t&inject.Object{Value: &replicaSet},\n\t\t&inject.Object{Value: &stats.HookClient{}},\n\t)\n\tensure.Nil(t, err)\n\tensure.Nil(t, graph.Populate())\n\tobjects := graph.Objects()\n\tensure.Nil(t, startstop.Start(objects, &log))\n\treturn &Harness{\n\t\tT:          t,\n\t\tStopper:    s,\n\t\tReplicaSet: &replicaSet,\n\t\tGraph:      &graph,\n\t\tLog:        &log,\n\t}\n}\n\ntype SingleHarness struct {\n\t*Harness\n\tMgoServer *mgotest.Server\n}\n\nfunc NewSingleHarness(t testing.TB) *SingleHarness {\n\tmgoserver := mgotest.NewStartedServer(t)\n\treturn &SingleHarness{\n\t\tHarness:   newHarnessInternal(mgoserver.URL(), mgoserver, t),\n\t\tMgoServer: mgoserver,\n\t}\n}\n\ntype ReplicaSetHarness struct {\n\t*Harness\n\tMgoReplicaSet *mgotest.ReplicaSet\n}\n\nfunc NewReplicaSetHarness(n uint, t testing.TB) *ReplicaSetHarness {\n\tif disableSlowTests {\n\t\tt.Skip(\"disabled because it's slow\")\n\t}\n\tmgoRS := mgotest.NewReplicaSet(n, t)\n\treturn &ReplicaSetHarness{\n\t\tHarness:       newHarnessInternal(mgoRS.Addrs()[n-1], mgoRS, t),\n\t\tMgoReplicaSet: mgoRS,\n\t}\n}\n\nfunc (h *Harness) Stop() {\n\tdefer h.Stopper.Stop()\n\tensure.Nil(h.T, startstop.Stop(h.Graph.Objects(), h.Log))\n}\n\nfunc (h *Harness) ProxySession() *mgo.Session {\n\treturn h.Dial(h.ReplicaSet.ProxyMembers()[0])\n}\n\nfunc (h *Harness) RealSession() *mgo.Session {\n\treturn h.Dial(h.ReplicaSet.lastState.Addrs()[0])\n}\n\nfunc (h *Harness) Dial(u string) *mgo.Session {\n\tsession, err := mgo.Dial(u)\n\tensure.Nil(h.T, err, u)\n\tsession.SetSafe(&mgo.Safe{FSync: true, W: 1})\n\tsession.SetSyncTimeout(time.Minute)\n\tsession.SetSocketTimeout(time.Minute)\n\treturn session\n}\n"
  },
  {
    "path": "doc.go",
    "content": "// Package dvara provides a library to enable setting up a proxy server\n// for mongo.\npackage dvara\n"
  },
  {
    "path": "license",
    "content": "BSD License\n\nFor dvara software\n\nCopyright (c) 2015, Facebook, Inc. All rights reserved.\n\nRedistribution and use in source and binary forms, with or without modification,\nare permitted provided that the following conditions are met:\n\n * Redistributions of source code must retain the above copyright notice, this\n   list of conditions and the following disclaimer.\n\n * Redistributions in binary form must reproduce the above copyright notice,\n   this list of conditions and the following disclaimer in the documentation\n   and/or other materials provided with the distribution.\n\n * Neither the name Facebook nor the names of its contributors may be used to\n   endorse or promote products derived from this software without specific\n   prior written permission.\n\nTHIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\nANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\nWARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\nDISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR\nANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\nLOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON\nANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\nSOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n"
  },
  {
    "path": "patents",
    "content": "Additional Grant of Patent Rights Version 2\n\n\"Software\" means the dvara software distributed by Facebook, Inc.\n\nFacebook, Inc. (\"Facebook\") hereby grants to each recipient of the Software\n(\"you\") a perpetual, worldwide, royalty-free, non-exclusive, irrevocable\n(subject to the termination provision below) license under any Necessary\nClaims, to make, have made, use, sell, offer to sell, import, and otherwise\ntransfer the Software. For avoidance of doubt, no license is granted under\nFacebook’s rights in any patent claims that are infringed by (i) modifications\nto the Software made by you or any third party or (ii) the Software in\ncombination with any software or other technology.\n\nThe license granted hereunder will terminate, automatically and without notice,\nif you (or any of your subsidiaries, corporate affiliates or agents) initiate\ndirectly or indirectly, or take a direct financial interest in, any Patent\nAssertion: (i) against Facebook or any of its subsidiaries or corporate\naffiliates, (ii) against any party if such Patent Assertion arises in whole or\nin part from any software, technology, product or service of Facebook or any of\nits subsidiaries or corporate affiliates, or (iii) against any party relating\nto the Software. Notwithstanding the foregoing, if Facebook or any of its\nsubsidiaries or corporate affiliates files a lawsuit alleging patent\ninfringement against you in the first instance, and you respond by filing a\npatent infringement counterclaim in that lawsuit against that party that is\nunrelated to the Software, the license granted hereunder will not terminate\nunder section (i) of this paragraph due to such counterclaim.\n\nA \"Necessary Claim\" is a claim of a patent owned by Facebook that is\nnecessarily infringed by the Software standing alone.\n\nA \"Patent Assertion\" is any lawsuit or other action alleging direct, indirect,\nor contributory infringement or inducement to infringe any patent, including a\ncross-claim or counterclaim.\n"
  },
  {
    "path": "protocol.go",
    "content": "package dvara\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n)\n\nvar (\n\terrWrite = errors.New(\"incorrect number of bytes written\")\n)\n\n// Look at http://docs.mongodb.org/meta-driver/latest/legacy/mongodb-wire-protocol/ for the protocol.\n\n// OpCode allow identifying the type of operation:\n//\n// http://docs.mongodb.org/meta-driver/latest/legacy/mongodb-wire-protocol/#request-opcodes\ntype OpCode int32\n\n// String returns a human readable representation of the OpCode.\nfunc (c OpCode) String() string {\n\tswitch c {\n\tdefault:\n\t\treturn \"UNKNOWN\"\n\tcase OpReply:\n\t\treturn \"REPLY\"\n\tcase OpMessage:\n\t\treturn \"MESSAGE\"\n\tcase OpUpdate:\n\t\treturn \"UPDATE\"\n\tcase OpInsert:\n\t\treturn \"INSERT\"\n\tcase Reserved:\n\t\treturn \"RESERVED\"\n\tcase OpQuery:\n\t\treturn \"QUERY\"\n\tcase OpGetMore:\n\t\treturn \"GET_MORE\"\n\tcase OpDelete:\n\t\treturn \"DELETE\"\n\tcase OpKillCursors:\n\t\treturn \"KILL_CURSORS\"\n\t}\n}\n\n// IsMutation tells us if the operation will mutate data. These operations can\n// be followed up by a getLastErr operation.\nfunc (c OpCode) IsMutation() bool {\n\treturn c == OpInsert || c == OpUpdate || c == OpDelete\n}\n\n// HasResponse tells us if the operation will have a response from the server.\nfunc (c OpCode) HasResponse() bool {\n\treturn c == OpQuery || c == OpGetMore\n}\n\n// The full set of known request op codes:\n// http://docs.mongodb.org/meta-driver/latest/legacy/mongodb-wire-protocol/#request-opcodes\nconst (\n\tOpReply       = OpCode(1)\n\tOpMessage     = OpCode(1000)\n\tOpUpdate      = OpCode(2001)\n\tOpInsert      = OpCode(2002)\n\tReserved      = OpCode(2003)\n\tOpQuery       = OpCode(2004)\n\tOpGetMore     = OpCode(2005)\n\tOpDelete      = OpCode(2006)\n\tOpKillCursors = OpCode(2007)\n)\n\n// messageHeader is the mongo MessageHeader\ntype messageHeader struct {\n\t// MessageLength is the total message size, including this header\n\tMessageLength int32\n\t// RequestID is the identifier for this miessage\n\tRequestID int32\n\t// ResponseTo is the RequestID of the message being responded to. used in DB responses\n\tResponseTo int32\n\t// OpCode is the request type, see consts above.\n\tOpCode OpCode\n}\n\n// ToWire converts the messageHeader to the wire protocol\nfunc (m messageHeader) ToWire() []byte {\n\tvar d [headerLen]byte\n\tb := d[:]\n\tsetInt32(b, 0, m.MessageLength)\n\tsetInt32(b, 4, m.RequestID)\n\tsetInt32(b, 8, m.ResponseTo)\n\tsetInt32(b, 12, int32(m.OpCode))\n\treturn b\n}\n\n// FromWire reads the wirebytes into this object\nfunc (m *messageHeader) FromWire(b []byte) {\n\tm.MessageLength = getInt32(b, 0)\n\tm.RequestID = getInt32(b, 4)\n\tm.ResponseTo = getInt32(b, 8)\n\tm.OpCode = OpCode(getInt32(b, 12))\n}\n\nfunc (m *messageHeader) WriteTo(w io.Writer) error {\n\tb := m.ToWire()\n\tn, err := w.Write(b)\n\tif err != nil {\n\t\treturn err\n\t}\n\tif n != len(b) {\n\t\treturn errWrite\n\t}\n\treturn nil\n}\n\n// String returns a string representation of the message header. Useful for debugging.\nfunc (m *messageHeader) String() string {\n\treturn fmt.Sprintf(\n\t\t\"opCode:%s (%d) msgLen:%d reqID:%d respID:%d\",\n\t\tm.OpCode,\n\t\tm.OpCode,\n\t\tm.MessageLength,\n\t\tm.RequestID,\n\t\tm.ResponseTo,\n\t)\n}\n\nfunc readHeader(r io.Reader) (*messageHeader, error) {\n\tvar d [headerLen]byte\n\tb := d[:]\n\tif _, err := io.ReadFull(r, b); err != nil {\n\t\treturn nil, err\n\t}\n\th := messageHeader{}\n\th.FromWire(b)\n\treturn &h, nil\n}\n\n// copyMessage copies reads & writes an entire message.\nfunc copyMessage(w io.Writer, r io.Reader) error {\n\th, err := readHeader(r)\n\tif err != nil {\n\t\treturn err\n\t}\n\tif err := h.WriteTo(w); err != nil {\n\t\treturn err\n\t}\n\t_, err = io.CopyN(w, r, int64(h.MessageLength-headerLen))\n\treturn err\n}\n\n// readDocument read an entire BSON document. This document can be used with\n// bson.Unmarshal.\nfunc readDocument(r io.Reader) ([]byte, error) {\n\tvar sizeRaw [4]byte\n\tif _, err := io.ReadFull(r, sizeRaw[:]); err != nil {\n\t\treturn nil, err\n\t}\n\tsize := getInt32(sizeRaw[:], 0)\n\tdoc := make([]byte, size)\n\tsetInt32(doc, 0, size)\n\tif _, err := io.ReadFull(r, doc[4:]); err != nil {\n\t\treturn nil, err\n\t}\n\treturn doc, nil\n}\n\nconst x00 = byte(0)\n\n// readCString reads a null turminated string as defined by BSON from the\n// reader. Note, the return value includes the trailing null byte.\nfunc readCString(r io.Reader) ([]byte, error) {\n\tvar b []byte\n\tvar n [1]byte\n\tfor {\n\t\tif _, err := io.ReadFull(r, n[:]); err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tb = append(b, n[0])\n\t\tif n[0] == x00 {\n\t\t\treturn b, nil\n\t\t}\n\t}\n}\n\n// all data in the MongoDB wire protocol is little-endian.\n// all the read/write functions below are little-endian.\nfunc getInt32(b []byte, pos int) int32 {\n\treturn (int32(b[pos+0])) |\n\t\t(int32(b[pos+1]) << 8) |\n\t\t(int32(b[pos+2]) << 16) |\n\t\t(int32(b[pos+3]) << 24)\n}\n\nfunc setInt32(b []byte, pos int, i int32) {\n\tb[pos] = byte(i)\n\tb[pos+1] = byte(i >> 8)\n\tb[pos+2] = byte(i >> 16)\n\tb[pos+3] = byte(i >> 24)\n}\n"
  },
  {
    "path": "protocol_test.go",
    "content": "package dvara\n\nimport (\n\t\"bytes\"\n\t\"errors\"\n\t\"io\"\n\t\"testing\"\n)\n\ntype testReader struct {\n\tread func([]byte) (int, error)\n}\n\nfunc (t testReader) Read(b []byte) (int, error) { return t.read(b) }\n\ntype testWriter struct {\n\twrite func([]byte) (int, error)\n}\n\nfunc (t testWriter) Write(b []byte) (int, error) { return t.write(b) }\n\nfunc TestOpStrings(t *testing.T) {\n\tt.Parallel()\n\tcases := []struct {\n\t\tOpCode OpCode\n\t\tString string\n\t}{\n\t\t{OpCode(0), \"UNKNOWN\"},\n\t\t{OpReply, \"REPLY\"},\n\t\t{OpMessage, \"MESSAGE\"},\n\t\t{OpUpdate, \"UPDATE\"},\n\t\t{OpInsert, \"INSERT\"},\n\t\t{Reserved, \"RESERVED\"},\n\t\t{OpQuery, \"QUERY\"},\n\t\t{OpGetMore, \"GET_MORE\"},\n\t\t{OpDelete, \"DELETE\"},\n\t\t{OpKillCursors, \"KILL_CURSORS\"},\n\t}\n\tfor _, c := range cases {\n\t\tif c.OpCode.String() != c.String {\n\t\t\tt.Fatalf(\"for code %d expected %s but got %s\", c.OpCode, c.String, c.OpCode)\n\t\t}\n\t}\n}\n\nfunc TestMsgHeaderString(t *testing.T) {\n\tt.Parallel()\n\tm := &messageHeader{\n\t\tOpCode:        OpQuery,\n\t\tMessageLength: 10,\n\t\tRequestID:     42,\n\t\tResponseTo:    43,\n\t}\n\tif m.String() != \"opCode:QUERY (2004) msgLen:10 reqID:42 respID:43\" {\n\t\tt.Fatalf(\"did not find expected string, instead found: %s\", m)\n\t}\n}\n\nfunc TestCopyEmptyMessage(t *testing.T) {\n\tt.Parallel()\n\tmsg := messageHeader{}\n\tmsgBytes := msg.ToWire()\n\tr := bytes.NewReader(msgBytes)\n\tvar w bytes.Buffer\n\tif err := copyMessage(&w, r); err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif !bytes.Equal(msgBytes, w.Bytes()) {\n\t\tt.Fatalf(\"did not get expected bytes %v got %v\", msgBytes, w.Bytes())\n\t}\n}\n\nfunc TestCopyMessageFromReadError(t *testing.T) {\n\tt.Parallel()\n\texpectedErr := errors.New(\"foo\")\n\tr := testReader{\n\t\tread: func(b []byte) (int, error) {\n\t\t\treturn 0, expectedErr\n\t\t},\n\t}\n\tvar w bytes.Buffer\n\tif err := copyMessage(&w, r); err != expectedErr {\n\t\tt.Fatalf(\"did not get expected error, instead got: %s\", err)\n\t}\n}\n\nfunc TestCopyMessageFromWriteError(t *testing.T) {\n\tt.Parallel()\n\tmsg := messageHeader{}\n\tr := bytes.NewReader(msg.ToWire())\n\texpectedErr := errors.New(\"foo\")\n\tw := testWriter{\n\t\twrite: func(b []byte) (int, error) {\n\t\t\treturn 0, expectedErr\n\t\t},\n\t}\n\tif err := copyMessage(w, r); err != expectedErr {\n\t\tt.Fatalf(\"did not get expected error, instead got: %s\", err)\n\t}\n}\n\nfunc TestCopyMessageFromWriteLengthError(t *testing.T) {\n\tt.Parallel()\n\tmsg := messageHeader{}\n\tr := bytes.NewReader(msg.ToWire())\n\tw := testWriter{\n\t\twrite: func(b []byte) (int, error) {\n\t\t\treturn 0, nil\n\t\t},\n\t}\n\tif err := copyMessage(w, r); err != errWrite {\n\t\tt.Fatalf(\"did not get expected error, instead got: %s\", err)\n\t}\n}\n\nfunc TestReadDocumentEmpty(t *testing.T) {\n\tt.Parallel()\n\tdoc, err := readDocument(bytes.NewReader([]byte{}))\n\tif err != io.EOF {\n\t\tt.Fatal(\"did not find expected error\")\n\t}\n\tif len(doc) != 0 {\n\t\tt.Fatal(\"was expecting an empty document\")\n\t}\n}\n\nfunc TestReadDocumentPartial(t *testing.T) {\n\tt.Parallel()\n\tfirst := true\n\tr := testReader{\n\t\tread: func(b []byte) (int, error) {\n\t\t\tif first {\n\t\t\t\tfirst = false\n\t\t\t\tsetInt32(b, 0, 5)\n\t\t\t\treturn 4, nil\n\t\t\t}\n\t\t\treturn 0, io.EOF\n\t\t},\n\t}\n\tdoc, err := readDocument(r)\n\tif err != io.EOF {\n\t\tt.Fatalf(\"did not find expected error, instead got %s %v\", err, doc)\n\t}\n\tif len(doc) != 0 {\n\t\tt.Fatal(\"was expecting an empty document\")\n\t}\n}\n\nfunc TestReadCString(t *testing.T) {\n\tt.Parallel()\n\tcases := []struct {\n\t\tData     []byte\n\t\tExpected []byte\n\t\tError    error\n\t}{\n\t\t{nil, nil, io.EOF},\n\t\t{[]byte{0}, []byte{0}, nil},\n\t\t{[]byte{1, 2, 3, 0}, []byte{1, 2, 3, 0}, nil},\n\t\t{[]byte{1, 0, 3}, []byte{1, 0}, nil},\n\t}\n\n\tfor _, c := range cases {\n\t\tcstring, err := readCString(bytes.NewReader(c.Data))\n\t\tif err != c.Error {\n\t\t\tt.Fatalf(\"did not find expected error, instead got %s %v\", err, cstring)\n\t\t}\n\t\tif !bytes.Equal(c.Expected, cstring) {\n\t\t\tt.Fatalf(\"did not find expected %v instead got %v\", c.Expected, cstring)\n\t\t}\n\t}\n}\n"
  },
  {
    "path": "proxy.go",
    "content": "package dvara\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"net\"\n\t\"os\"\n\t\"strings\"\n\t\"sync\"\n\t\"time\"\n\n\t\"github.com/facebookgo/rpool\"\n\t\"github.com/facebookgo/stats\"\n)\n\nconst headerLen = 16\n\nvar (\n\terrZeroMaxConnections          = errors.New(\"dvara: MaxConnections cannot be 0\")\n\terrZeroMaxPerClientConnections = errors.New(\"dvara: MaxPerClientConnections cannot be 0\")\n\terrNormalClose                 = errors.New(\"dvara: normal close\")\n\terrClientReadTimeout           = errors.New(\"dvara: client read timeout\")\n\n\ttimeInPast = time.Now()\n)\n\n// Proxy sends stuff from clients to mongo servers.\ntype Proxy struct {\n\tLog            Logger\n\tReplicaSet     *ReplicaSet\n\tClientListener net.Listener // Listener for incoming client connections\n\tProxyAddr      string       // Address for incoming client connections\n\tMongoAddr      string       // Address for destination Mongo server\n\n\twg                      sync.WaitGroup\n\tclosed                  chan struct{}\n\tserverPool              rpool.Pool\n\tstats                   stats.Client\n\tmaxPerClientConnections *maxPerClientConnections\n}\n\n// String representation for debugging.\nfunc (p *Proxy) String() string {\n\treturn fmt.Sprintf(\"proxy %s => mongo %s\", p.ProxyAddr, p.MongoAddr)\n}\n\n// Start the proxy.\nfunc (p *Proxy) Start() error {\n\tif p.ReplicaSet.MaxConnections == 0 {\n\t\treturn errZeroMaxConnections\n\t}\n\tif p.ReplicaSet.MaxPerClientConnections == 0 {\n\t\treturn errZeroMaxPerClientConnections\n\t}\n\n\tp.closed = make(chan struct{})\n\tp.maxPerClientConnections = newMaxPerClientConnections(p.ReplicaSet.MaxPerClientConnections)\n\tp.serverPool = rpool.Pool{\n\t\tNew:               p.newServerConn,\n\t\tCloseErrorHandler: p.serverCloseErrorHandler,\n\t\tMax:               p.ReplicaSet.MaxConnections,\n\t\tMinIdle:           p.ReplicaSet.MinIdleConnections,\n\t\tIdleTimeout:       p.ReplicaSet.ServerIdleTimeout,\n\t\tClosePoolSize:     p.ReplicaSet.ServerClosePoolSize,\n\t}\n\n\t// plug stats if we can\n\tif p.ReplicaSet.Stats != nil {\n\t\t// Drop the default port suffix to make them pretty in production.\n\t\tdbName := strings.TrimSuffix(p.MongoAddr, \":27017\")\n\n\t\t// We want 2 sets of keys, one specific to the proxy, and another shared\n\t\t// with others.\n\t\tp.serverPool.Stats = stats.PrefixClient(\n\t\t\t[]string{\n\t\t\t\t\"mongoproxy.server.pool.\",\n\t\t\t\tfmt.Sprintf(\"mongoproxy.%s.server.pool.\", dbName),\n\t\t\t},\n\t\t\tp.ReplicaSet.Stats,\n\t\t)\n\t\tp.stats = stats.PrefixClient(\n\t\t\t[]string{\n\t\t\t\t\"mongoproxy.\",\n\t\t\t\tfmt.Sprintf(\"mongoproxy.%s.\", dbName),\n\t\t\t},\n\t\t\tp.ReplicaSet.Stats,\n\t\t)\n\t}\n\n\tgo p.clientAcceptLoop()\n\n\treturn nil\n}\n\n// Stop the proxy.\nfunc (p *Proxy) Stop() error {\n\treturn p.stop(false)\n}\n\nfunc (p *Proxy) stop(hard bool) error {\n\tif err := p.ClientListener.Close(); err != nil {\n\t\treturn err\n\t}\n\tclose(p.closed)\n\tif !hard {\n\t\tp.wg.Wait()\n\t}\n\tp.serverPool.Close()\n\treturn nil\n}\n\nfunc (p *Proxy) checkRSChanged() bool {\n\taddrs := p.ReplicaSet.lastState.Addrs()\n\tr, err := p.ReplicaSet.ReplicaSetStateCreator.FromAddrs(addrs, p.ReplicaSet.Name)\n\tif err != nil {\n\t\tp.Log.Errorf(\"all nodes possibly down?: %s\", err)\n\t\treturn true\n\t}\n\n\tif err := r.AssertEqual(p.ReplicaSet.lastState); err != nil {\n\t\tp.Log.Error(err)\n\t\tgo p.ReplicaSet.Restart()\n\t\treturn true\n\t}\n\n\treturn false\n}\n\n// Open up a new connection to the server. Retry 7 times, doubling the sleep\n// each time. This means we'll a total of 12.75 seconds with the last wait\n// being 6.4 seconds.\nfunc (p *Proxy) newServerConn() (io.Closer, error) {\n\tretrySleep := 50 * time.Millisecond\n\tfor retryCount := 7; retryCount > 0; retryCount-- {\n\t\tc, err := net.Dial(\"tcp\", p.MongoAddr)\n\t\tif err == nil {\n\t\t\treturn c, nil\n\t\t}\n\t\tp.Log.Error(err)\n\n\t\t// abort if rs changed\n\t\tif p.checkRSChanged() {\n\t\t\treturn nil, errNormalClose\n\t\t}\n\t\ttime.Sleep(retrySleep)\n\t\tretrySleep = retrySleep * 2\n\t}\n\treturn nil, fmt.Errorf(\"could not connect to %s\", p.MongoAddr)\n}\n\n// getServerConn gets a server connection from the pool.\nfunc (p *Proxy) getServerConn() (net.Conn, error) {\n\tc, err := p.serverPool.Acquire()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn c.(net.Conn), nil\n}\n\nfunc (p *Proxy) serverCloseErrorHandler(err error) {\n\tp.Log.Error(err)\n}\n\n// proxyMessage proxies a message, possibly it's response, and possibly a\n// follow up call.\nfunc (p *Proxy) proxyMessage(\n\th *messageHeader,\n\tclient net.Conn,\n\tserver net.Conn,\n\tlastError *LastError,\n) error {\n\n\tp.Log.Debugf(\"proxying message %s from %s for %s\", h, client.RemoteAddr(), p)\n\tdeadline := time.Now().Add(p.ReplicaSet.MessageTimeout)\n\tserver.SetDeadline(deadline)\n\tclient.SetDeadline(deadline)\n\n\t// OpQuery may need to be transformed and need special handling in order to\n\t// make the proxy transparent.\n\tif h.OpCode == OpQuery {\n\t\tstats.BumpSum(p.stats, \"message.with.response\", 1)\n\t\treturn p.ReplicaSet.ProxyQuery.Proxy(h, client, server, lastError)\n\t}\n\n\t// Anything besides a getlasterror call (which requires an OpQuery) resets\n\t// the lastError.\n\tif lastError.Exists() {\n\t\tp.Log.Debug(\"reset getLastError cache\")\n\t\tlastError.Reset()\n\t}\n\n\t// For other Ops we proxy the header & raw body over.\n\tif err := h.WriteTo(server); err != nil {\n\t\tp.Log.Error(err)\n\t\treturn err\n\t}\n\n\tif _, err := io.CopyN(server, client, int64(h.MessageLength-headerLen)); err != nil {\n\t\tp.Log.Error(err)\n\t\treturn err\n\t}\n\n\t// For Ops with responses we proxy the raw response message over.\n\tif h.OpCode.HasResponse() {\n\t\tstats.BumpSum(p.stats, \"message.with.response\", 1)\n\t\tif err := copyMessage(client, server); err != nil {\n\t\t\tp.Log.Error(err)\n\t\t\treturn err\n\t\t}\n\t}\n\n\treturn nil\n}\n\n// clientAcceptLoop accepts new clients and creates a clientServeLoop for each\n// new client that connects to the proxy.\nfunc (p *Proxy) clientAcceptLoop() {\n\tfor {\n\t\tp.wg.Add(1)\n\t\tc, err := p.ClientListener.Accept()\n\t\tif err != nil {\n\t\t\tp.wg.Done()\n\t\t\tif strings.Contains(err.Error(), \"use of closed network connection\") {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tp.Log.Error(err)\n\t\t\tcontinue\n\t\t}\n\t\tgo p.clientServeLoop(c)\n\t}\n}\n\n// clientServeLoop loops on a single client connected to the proxy and\n// dispatches its requests.\nfunc (p *Proxy) clientServeLoop(c net.Conn) {\n\tremoteIP := c.RemoteAddr().(*net.TCPAddr).IP.String()\n\n\t// enforce per-client max connection limit\n\tif p.maxPerClientConnections.inc(remoteIP) {\n\t\tc.Close()\n\t\tstats.BumpSum(p.stats, \"client.rejected.max.connections\", 1)\n\t\tp.Log.Errorf(\"rejecting client connection due to max connections limit: %s\", remoteIP)\n\t\treturn\n\t}\n\n\t// turn on TCP keep-alive and set it to the recommended period of 2 minutes\n\t// http://docs.mongodb.org/manual/faq/diagnostics/#faq-keepalive\n\tif conn, ok := c.(*net.TCPConn); ok {\n\t\tconn.SetKeepAlivePeriod(2 * time.Minute)\n\t\tconn.SetKeepAlive(true)\n\t}\n\n\tc = teeIf(fmt.Sprintf(\"client %s <=> %s\", c.RemoteAddr(), p), c)\n\tp.Log.Infof(\"client %s connected to %s\", c.RemoteAddr(), p)\n\tstats.BumpSum(p.stats, \"client.connected\", 1)\n\tdefer func() {\n\t\tp.Log.Infof(\"client %s disconnected from %s\", c.RemoteAddr(), p)\n\t\tp.wg.Done()\n\t\tif err := c.Close(); err != nil {\n\t\t\tp.Log.Error(err)\n\t\t}\n\t\tp.maxPerClientConnections.dec(remoteIP)\n\t}()\n\n\tvar lastError LastError\n\tfor {\n\t\th, err := p.idleClientReadHeader(c)\n\t\tif err != nil {\n\t\t\tif err != errNormalClose {\n\t\t\t\tp.Log.Error(err)\n\t\t\t}\n\t\t\treturn\n\t\t}\n\n\t\tmpt := stats.BumpTime(p.stats, \"message.proxy.time\")\n\t\tserverConn, err := p.getServerConn()\n\t\tif err != nil {\n\t\t\tif err != errNormalClose {\n\t\t\t\tp.Log.Error(err)\n\t\t\t}\n\t\t\treturn\n\t\t}\n\n\t\tscht := stats.BumpTime(p.stats, \"server.conn.held.time\")\n\t\tfor {\n\t\t\terr := p.proxyMessage(h, c, serverConn, &lastError)\n\t\t\tif err != nil {\n\t\t\t\tp.serverPool.Discard(serverConn)\n\t\t\t\tp.Log.Error(err)\n\t\t\t\tstats.BumpSum(p.stats, \"message.proxy.error\", 1)\n\t\t\t\tif ne, ok := err.(net.Error); ok && ne.Timeout() {\n\t\t\t\t\tstats.BumpSum(p.stats, \"message.proxy.timeout\", 1)\n\t\t\t\t}\n\t\t\t\tif err == errRSChanged {\n\t\t\t\t\tgo p.ReplicaSet.Restart()\n\t\t\t\t}\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\t// One message was proxied, stop it's timer.\n\t\t\tmpt.End()\n\n\t\t\tif !h.OpCode.IsMutation() {\n\t\t\t\tbreak\n\t\t\t}\n\n\t\t\t// If the operation we just performed was a mutation, we always make the\n\t\t\t// follow up request on the same server because it's possibly a getLastErr\n\t\t\t// call which expects this behavior.\n\n\t\t\tstats.BumpSum(p.stats, \"message.with.mutation\", 1)\n\t\t\th, err = p.gleClientReadHeader(c)\n\t\t\tif err != nil {\n\t\t\t\t// Client did not make _any_ query within the GetLastErrorTimeout.\n\t\t\t\t// Return the server to the pool and wait go back to outer loop.\n\t\t\t\tif err == errClientReadTimeout {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t\t// Prevent noise of normal client disconnects, but log if anything else.\n\t\t\t\tif err != errNormalClose {\n\t\t\t\t\tp.Log.Error(err)\n\t\t\t\t}\n\t\t\t\t// We need to return our server to the pool (it's still good as far\n\t\t\t\t// as we know).\n\t\t\t\tp.serverPool.Release(serverConn)\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\t// Successfully read message when waiting for the getLastError call.\n\t\t\tmpt = stats.BumpTime(p.stats, \"message.proxy.time\")\n\t\t}\n\t\tp.serverPool.Release(serverConn)\n\t\tscht.End()\n\t\tstats.BumpSum(p.stats, \"message.proxy.success\", 1)\n\t}\n}\n\n// We wait for upto ClientIdleTimeout in MessageTimeout increments and keep\n// checking if we're waiting to be closed. This ensures that at worse we\n// wait for MessageTimeout when closing even when we're idling.\nfunc (p *Proxy) idleClientReadHeader(c net.Conn) (*messageHeader, error) {\n\th, err := p.clientReadHeader(c, p.ReplicaSet.ClientIdleTimeout)\n\tif err == errClientReadTimeout {\n\t\tstats.BumpSum(p.stats, \"client.idle.timeout\", 1)\n\t}\n\treturn h, err\n}\n\nfunc (p *Proxy) gleClientReadHeader(c net.Conn) (*messageHeader, error) {\n\th, err := p.clientReadHeader(c, p.ReplicaSet.GetLastErrorTimeout)\n\tif err == errClientReadTimeout {\n\t\tstats.BumpSum(p.stats, \"client.gle.timeout\", 1)\n\t}\n\treturn h, err\n}\n\nfunc (p *Proxy) clientReadHeader(c net.Conn, timeout time.Duration) (*messageHeader, error) {\n\tt := stats.BumpTime(p.stats, \"client.read.header.time\")\n\ttype headerError struct {\n\t\theader *messageHeader\n\t\terror  error\n\t}\n\tresChan := make(chan headerError)\n\n\tc.SetReadDeadline(time.Now().Add(timeout))\n\tgo func() {\n\t\th, err := readHeader(c)\n\t\tresChan <- headerError{header: h, error: err}\n\t}()\n\n\tclosed := false\n\tvar response headerError\n\n\tselect {\n\tcase response = <-resChan:\n\t\t// all good\n\tcase <-p.closed:\n\t\tclosed = true\n\t\tc.SetReadDeadline(timeInPast)\n\t\tresponse = <-resChan\n\t}\n\n\t// Successfully read a header.\n\tif response.error == nil {\n\t\tt.End()\n\t\treturn response.header, nil\n\t}\n\n\t// Client side disconnected.\n\tif response.error == io.EOF {\n\t\tstats.BumpSum(p.stats, \"client.clean.disconnect\", 1)\n\t\treturn nil, errNormalClose\n\t}\n\n\t// We hit our ReadDeadline.\n\tif ne, ok := response.error.(net.Error); ok && ne.Timeout() {\n\t\tif closed {\n\t\t\tstats.BumpSum(p.stats, \"client.clean.disconnect\", 1)\n\t\t\treturn nil, errNormalClose\n\t\t}\n\t\treturn nil, errClientReadTimeout\n\t}\n\n\t// Some other unknown error.\n\tstats.BumpSum(p.stats, \"client.error.disconnect\", 1)\n\tp.Log.Error(response.error)\n\treturn nil, response.error\n}\n\nvar teeIfEnable = os.Getenv(\"MONGOPROXY_TEE\") == \"1\"\n\ntype teeConn struct {\n\tcontext string\n\tnet.Conn\n}\n\nfunc (t teeConn) Read(b []byte) (int, error) {\n\tn, err := t.Conn.Read(b)\n\tif n > 0 {\n\t\tfmt.Fprintf(os.Stdout, \"READ %s: %s %v\\n\", t.context, b[0:n], b[0:n])\n\t}\n\treturn n, err\n}\n\nfunc (t teeConn) Write(b []byte) (int, error) {\n\tn, err := t.Conn.Write(b)\n\tif n > 0 {\n\t\tfmt.Fprintf(os.Stdout, \"WRIT %s: %s %v\\n\", t.context, b[0:n], b[0:n])\n\t}\n\treturn n, err\n}\n\nfunc teeIf(context string, c net.Conn) net.Conn {\n\tif teeIfEnable {\n\t\treturn teeConn{\n\t\t\tcontext: context,\n\t\t\tConn:    c,\n\t\t}\n\t}\n\treturn c\n}\n\ntype maxPerClientConnections struct {\n\tmax    uint\n\tcounts map[string]uint\n\tmutex  sync.Mutex\n}\n\nfunc newMaxPerClientConnections(max uint) *maxPerClientConnections {\n\treturn &maxPerClientConnections{\n\t\tmax:    max,\n\t\tcounts: make(map[string]uint),\n\t}\n}\n\nfunc (m *maxPerClientConnections) inc(remoteIP string) bool {\n\tm.mutex.Lock()\n\tdefer m.mutex.Unlock()\n\tcurrent := m.counts[remoteIP]\n\tif current >= m.max {\n\t\treturn true\n\t}\n\tm.counts[remoteIP] = current + 1\n\treturn false\n}\n\nfunc (m *maxPerClientConnections) dec(remoteIP string) {\n\tm.mutex.Lock()\n\tdefer m.mutex.Unlock()\n\tcurrent := m.counts[remoteIP]\n\n\t// delete rather than having entries with 0 connections\n\tif current == 1 {\n\t\tdelete(m.counts, remoteIP)\n\t} else {\n\t\tm.counts[remoteIP] = current - 1\n\t}\n}\n"
  },
  {
    "path": "proxy_test.go",
    "content": "package dvara\n\nimport (\n\t\"fmt\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/facebookgo/ensure\"\n\t\"github.com/facebookgo/inject\"\n\t\"github.com/facebookgo/mgotest\"\n\t\"github.com/facebookgo/startstop\"\n\t\"github.com/facebookgo/stats\"\n\n\t\"gopkg.in/mgo.v2\"\n\t\"gopkg.in/mgo.v2/bson\"\n)\n\nfunc TestParallelInsertWithUniqueIndex(t *testing.T) {\n\tt.Parallel()\n\tif disableSlowTests {\n\t\tt.Skip(\"TestParallelInsertWithUniqueIndex disabled because it's slow\")\n\t}\n\th := NewSingleHarness(t)\n\tdefer h.Stop()\n\n\tlimit := 20000\n\tc := make(chan int, limit)\n\tfor i := 0; i < 3; i++ {\n\t\tgo inserter(h.ProxySession(), c, limit)\n\t}\n\tset := make(map[int]bool)\n\tfor k := range c {\n\t\tif set[k] {\n\t\t\tt.Fatal(\"Double write on same value\")\n\t\t}\n\t\tset[k] = true\n\t\tif len(set) == limit {\n\t\t\tbreak\n\t\t}\n\t}\n}\n\nfunc inserter(s *mgo.Session, channel chan int, limit int) {\n\tdefer s.Close()\n\tc := s.DB(\"test\").C(\"test\")\n\tc.EnsureIndex(mgo.Index{Key: []string{\"phoneNum\"}, Unique: true})\n\tfor i := 1; i <= limit; i++ {\n\t\tif err := c.Insert(bson.M{\"phoneNum\": i}); err == nil {\n\t\t\tchannel <- i\n\t\t}\n\t}\n}\nfunc TestSimpleCRUD(t *testing.T) {\n\tt.Parallel()\n\tp := NewSingleHarness(t)\n\tdefer p.Stop()\n\tsession := p.ProxySession()\n\tdefer session.Close()\n\tcollection := session.DB(\"test\").C(\"coll1\")\n\tdata := map[string]interface{}{\n\t\t\"_id\":  1,\n\t\t\"name\": \"abc\",\n\t}\n\terr := collection.Insert(data)\n\tif err != nil {\n\t\tt.Fatal(\"insertion error\", err)\n\t}\n\tn, err := collection.Count()\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif n != 1 {\n\t\tt.Fatalf(\"expecting 1 got %d\", n)\n\t}\n\tresult := make(map[string]interface{})\n\tcollection.Find(bson.M{\"_id\": 1}).One(&result)\n\tif result[\"name\"] != \"abc\" {\n\t\tt.Fatal(\"expecting name abc got\", result)\n\t}\n\terr = collection.DropCollection()\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n}\n\n// inserting data with same id field twice should fail\nfunc TestIDConstraint(t *testing.T) {\n\tt.Parallel()\n\tp := NewSingleHarness(t)\n\tdefer p.Stop()\n\tsession := p.ProxySession()\n\tdefer session.Close()\n\tcollection := session.DB(\"test\").C(\"coll1\")\n\tdata := map[string]interface{}{\n\t\t\"_id\":  1,\n\t\t\"name\": \"abc\",\n\t}\n\terr := collection.Insert(data)\n\tif err != nil {\n\t\tt.Fatal(\"insertion error\", err)\n\t}\n\terr = collection.Insert(data)\n\tif err == nil {\n\t\tt.Fatal(\"insertion failed on same id without write concern\")\n\t}\n}\n\n// inserting data voilating index clause on a separate connection should fail\nfunc TestEnsureIndex(t *testing.T) {\n\tt.Parallel()\n\tp := NewSingleHarness(t)\n\tdefer p.Stop()\n\tsession := p.ProxySession()\n\tcollection := session.DB(\"test\").C(\"coll1\")\n\tindex := mgo.Index{\n\t\tKey:        []string{\"lastname\", \"firstname\"},\n\t\tUnique:     true,\n\t\tDropDups:   true,\n\t\tBackground: true, // See notes.\n\t\tSparse:     true,\n\t}\n\terr := collection.EnsureIndex(index)\n\tensure.Nil(t, err)\n\terr = collection.Insert(\n\t\tmap[string]string{\n\t\t\t\"firstname\": \"harvey\",\n\t\t\t\"lastname\":  \"dent\",\n\t\t},\n\t)\n\tif err != nil {\n\t\tt.Fatal(\"insertion error\", err)\n\t}\n\tsession.Close()\n\tsession = p.ProxySession()\n\tdefer session.Close()\n\tcollection = session.DB(\"test\").C(\"coll1\")\n\terr = collection.Insert(\n\t\tmap[string]string{\n\t\t\t\"firstname\": \"harvey\",\n\t\t\t\"lastname\":  \"dent\",\n\t\t},\n\t)\n\tensure.NotNil(t, err)\n}\n\n// inserting same data after dropping an index should work\nfunc TestDropIndex(t *testing.T) {\n\tt.Parallel()\n\tp := NewSingleHarness(t)\n\tdefer p.Stop()\n\tsession := p.ProxySession()\n\tcollection := session.DB(\"test\").C(\"coll1\")\n\tindex := mgo.Index{\n\t\tKey:        []string{\"lastname\", \"firstname\"},\n\t\tUnique:     true,\n\t\tDropDups:   true,\n\t\tBackground: true, // See notes.\n\t\tSparse:     true,\n\t}\n\terr := collection.EnsureIndex(index)\n\tif err != nil {\n\t\tt.Fatal(\"ensure index call failed\")\n\t}\n\terr = collection.Insert(\n\t\tmap[string]string{\n\t\t\t\"firstname\": \"harvey\",\n\t\t\t\"lastname\":  \"dent\",\n\t\t},\n\t)\n\tif err != nil {\n\t\tt.Fatal(\"insertion error\", err)\n\t}\n\tcollection.DropIndex(\"lastname\", \"firstname\")\n\tsession.Close()\n\tsession = p.ProxySession()\n\tdefer session.Close()\n\tcollection = session.DB(\"test\").C(\"coll1\")\n\terr = collection.Insert(\n\t\tmap[string]string{\n\t\t\t\"firstname\": \"harvey\",\n\t\t\t\"lastname\":  \"dent\",\n\t\t},\n\t)\n\tif err != nil {\n\t\tt.Fatal(\"drop index did not work\")\n\t}\n}\n\nfunc TestRemoval(t *testing.T) {\n\tt.Parallel()\n\tp := NewSingleHarness(t)\n\tdefer p.Stop()\n\tsession := p.ProxySession()\n\tdefer session.Close()\n\tcollection := session.DB(\"test\").C(\"coll1\")\n\tif err := collection.Insert(bson.M{\"S\": \"hello\", \"I\": 24}); err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif err := collection.Remove(bson.M{\"S\": \"hello\", \"I\": 24}); err != nil {\n\t\tt.Fatal(err)\n\t}\n\tvar res []interface{}\n\tcollection.Find(bson.M{\"S\": \"hello\", \"I\": 24}).All(&res)\n\tif res != nil {\n\t\tt.Fatal(\"found object after delete\", res)\n\t}\n\tif err := collection.Remove(bson.M{\"S\": \"hello\", \"I\": 24}); err == nil {\n\t\tt.Fatal(\"removing nonexistant document should error\")\n\t}\n}\n\nfunc TestUpdate(t *testing.T) {\n\tt.Parallel()\n\tp := NewSingleHarness(t)\n\tdefer p.Stop()\n\tsession := p.ProxySession()\n\tdefer session.Close()\n\tcollection := session.DB(\"test\").C(\"coll1\")\n\tif err := collection.Insert(bson.M{\"_id\": \"1234\", \"name\": \"Alfred\"}); err != nil {\n\t\tt.Fatal(err)\n\t}\n\tvar result map[string]interface{}\n\tcollection.Find(nil).One(&result)\n\tif result[\"name\"] != \"Alfred\" {\n\t\tt.Fatal(\"insert failed\")\n\t}\n\tif err := collection.Update(bson.M{\"_id\": \"1234\"}, bson.M{\"name\": \"Jeeves\"}); err != nil {\n\t\tt.Fatal(\"update failed with\", err)\n\t}\n\tcollection.Find(nil).One(&result)\n\tif result[\"name\"] != \"Jeeves\" {\n\t\tt.Fatal(\"update failed\")\n\t}\n\tif err := collection.Update(bson.M{\"_id\": \"00000\"}, bson.M{\"name\": \"Jeeves\"}); err == nil {\n\t\tt.Fatal(\"update failed\")\n\t}\n}\n\nfunc TestStopChattyClient(t *testing.T) {\n\tt.Parallel()\n\tp := NewSingleHarness(t)\n\tsession := p.ProxySession()\n\tdefer session.Close()\n\tfin := make(chan struct{})\n\tgo func() {\n\t\tcollection := session.DB(\"test\").C(\"coll1\")\n\t\ti := 0\n\t\tfor {\n\t\t\tselect {\n\t\t\tdefault:\n\t\t\t\tcollection.Insert(bson.M{\"value\": i})\n\t\t\t\ti++\n\t\t\tcase <-fin:\n\t\t\t\treturn\n\t\t\t}\n\t\t}\n\t}()\n\tclose(fin)\n\tp.Stop()\n}\n\nfunc TestStopIdleClient(t *testing.T) {\n\tt.Parallel()\n\tp := NewSingleHarness(t)\n\tsession := p.ProxySession()\n\tdefer session.Close()\n\tif err := session.DB(\"test\").C(\"col\").Insert(bson.M{\"v\": 1}); err != nil {\n\t\tt.Fatal(err)\n\t}\n\tp.Stop()\n}\n\nfunc TestZeroMaxConnections(t *testing.T) {\n\tt.Parallel()\n\tp := &Proxy{ReplicaSet: &ReplicaSet{}}\n\terr := p.Start()\n\tif err != errZeroMaxConnections {\n\t\tt.Fatal(\"did not get expected error\")\n\t}\n}\n\nfunc TestNoAddrsGiven(t *testing.T) {\n\tt.Parallel()\n\treplicaSet := ReplicaSet{MaxConnections: 1}\n\tlog := tLogger{TB: t}\n\tvar graph inject.Graph\n\terr := graph.Provide(\n\t\t&inject.Object{Value: &log},\n\t\t&inject.Object{Value: &replicaSet},\n\t\t&inject.Object{Value: &stats.HookClient{}},\n\t)\n\tensure.Nil(t, err)\n\tensure.Nil(t, graph.Populate())\n\tobjects := graph.Objects()\n\terr = startstop.Start(objects, &log)\n\tif err != errNoAddrsGiven {\n\t\tt.Fatalf(\"did not get expected error, got: %s\", err)\n\t}\n}\n\nfunc TestSingleNodeWhenExpectingRS(t *testing.T) {\n\tt.Parallel()\n\tmgoserver := mgotest.NewStartedServer(t)\n\tdefer mgoserver.Stop()\n\treplicaSet := ReplicaSet{\n\t\tAddrs:          fmt.Sprintf(\"127.0.0.1:%d,127.0.0.1:%d\", mgoserver.Port, mgoserver.Port+1),\n\t\tMaxConnections: 1,\n\t}\n\tlog := tLogger{TB: t}\n\tvar graph inject.Graph\n\terr := graph.Provide(\n\t\t&inject.Object{Value: &log},\n\t\t&inject.Object{Value: &replicaSet},\n\t\t&inject.Object{Value: &stats.HookClient{}},\n\t)\n\tensure.Nil(t, err)\n\tensure.Nil(t, graph.Populate())\n\tobjects := graph.Objects()\n\terr = startstop.Start(objects, &log)\n\tif err == nil || !strings.Contains(err.Error(), \"was expecting it to be in a replica set\") {\n\t\tt.Fatalf(\"did not get expected error, got: %s\", err)\n\t}\n}\n\nfunc TestStopListenerCloseError(t *testing.T) {\n\tt.Parallel()\n\tp := NewSingleHarness(t)\n\tp.Stop()\n\terr := p.ReplicaSet.Stop()\n\tif err == nil || !strings.Contains(err.Error(), \"use of closed network connection\") {\n\t\tt.Fatalf(\"did not get expected error, instead got: %s\", err)\n\t}\n}\n\nfunc TestMongoGoingAwayAndReturning(t *testing.T) {\n\tt.Parallel()\n\tp := NewSingleHarness(t)\n\tsession := p.ProxySession()\n\tdefer session.Close()\n\tcollection := session.DB(\"test\").C(\"coll1\")\n\tif err := collection.Insert(bson.M{\"value\": 1}); err != nil {\n\t\tt.Fatal(err)\n\t}\n\tp.MgoServer.Stop()\n\tp.MgoServer.Start()\n\t// For now we can only gurantee that eventually things will work again. In an\n\t// ideal world the very first client connection after mongo returns should\n\t// work, and we shouldn't need a loop here.\n\tfor {\n\t\tcollection = session.Copy().DB(\"test\").C(\"coll1\")\n\t\tif err := collection.Insert(bson.M{\"value\": 3}); err == nil {\n\t\t\tbreak\n\t\t}\n\t}\n\tp.Stop()\n}\n\nfunc benchmarkInsertRead(b *testing.B, session *mgo.Session) {\n\tdefer session.Close()\n\tcol := session.DB(\"test\").C(\"col\")\n\tcol.EnsureIndex(mgo.Index{Key: []string{\"answer\"}, Unique: true})\n\tinsertDocs := bson.D{bson.DocElem{Name: \"answer\"}}\n\tinserted := bson.M{}\n\tb.ResetTimer()\n\tfor i := 0; i < b.N; i++ {\n\t\tinsertDocs[0].Value = i\n\t\tif err := col.Insert(insertDocs); err != nil {\n\t\t\tb.Fatal(err)\n\t\t}\n\t\tif err := col.Find(insertDocs).One(inserted); err != nil {\n\t\t\tb.Fatal(err)\n\t\t}\n\t\tif _, ok := inserted[\"_id\"]; !ok {\n\t\t\tb.Fatalf(\"no _id found: %+v\", inserted)\n\t\t}\n\t}\n}\n\nfunc BenchmarkInsertReadProxy(b *testing.B) {\n\tp := NewSingleHarness(b)\n\tbenchmarkInsertRead(b, p.ProxySession())\n}\n\nfunc BenchmarkInsertReadDirect(b *testing.B) {\n\tp := NewSingleHarness(b)\n\tbenchmarkInsertRead(b, p.RealSession())\n}\n"
  },
  {
    "path": "readme.md",
    "content": "dvara [![Build Status](https://secure.travis-ci.org/facebookgo/dvara.png)](http://travis-ci.org/facebookgo/dvara)\n=====\n\n**NOTE**: dvara is no longer in use and we are not accepting new pull requests\nfor it. Fork away and use it if it works for you!\n\n---\n\ndvara provides a connection pooling proxy for\n[MongoDB](http://www.mongodb.org/). For more information look at the associated\nblog post: http://blog.parse.com/2014/06/23/dvara/.\n\nTo build from source you'll need [Go](http://golang.org/). With it you can install it using:\n\n    go get github.com/facebookgo/dvara/cmd/dvara\n\nLibrary documentation: https://godoc.org/github.com/facebookgo/dvara\n"
  },
  {
    "path": "replica_set.go",
    "content": "package dvara\n\nimport (\n\t\"errors\"\n\t\"flag\"\n\t\"fmt\"\n\t\"net\"\n\t\"os\"\n\t\"strings\"\n\t\"sync\"\n\t\"time\"\n\n\t\"github.com/facebookgo/stackerr\"\n\t\"github.com/facebookgo/stats\"\n)\n\nvar hardRestart = flag.Bool(\n\t\"hard_restart\",\n\ttrue,\n\t\"if true will drop clients on restart\",\n)\n\n// Logger allows for simple text logging.\ntype Logger interface {\n\tError(args ...interface{})\n\tErrorf(format string, args ...interface{})\n\tWarn(args ...interface{})\n\tWarnf(format string, args ...interface{})\n\tInfo(args ...interface{})\n\tInfof(format string, args ...interface{})\n\tDebug(args ...interface{})\n\tDebugf(format string, args ...interface{})\n}\n\nvar errNoAddrsGiven = errors.New(\"dvara: no seed addresses given for ReplicaSet\")\n\n// ReplicaSet manages the real => proxy address mapping.\n// NewReplicaSet returns the ReplicaSet given the list of seed servers. It is\n// required for the seed servers to be a strict subset of the actual members if\n// they are reachable. That is, if two of the addresses are members of\n// different replica sets, it will be considered an error.\ntype ReplicaSet struct {\n\tLog                    Logger                  `inject:\"\"`\n\tReplicaSetStateCreator *ReplicaSetStateCreator `inject:\"\"`\n\tProxyQuery             *ProxyQuery             `inject:\"\"`\n\n\t// Stats if provided will be used to record interesting stats.\n\tStats stats.Client `inject:\"\"`\n\n\t// Comma separated list of mongo addresses. This is the list of \"seed\"\n\t// servers, and one of two conditions must be met for each entry here -- it's\n\t// either alive and part of the same replica set as all others listed, or is\n\t// not reachable.\n\tAddrs string\n\n\t// PortStart and PortEnd define the port range within which proxies will be\n\t// allocated.\n\tPortStart int\n\tPortEnd   int\n\n\t// Maximum number of connections that will be established to each mongo node.\n\tMaxConnections uint\n\n\t// MinIdleConnections is the number of idle server connections we'll keep\n\t// around.\n\tMinIdleConnections uint\n\n\t// ServerIdleTimeout is the duration after which a server connection will be\n\t// considered idle.\n\tServerIdleTimeout time.Duration\n\n\t// ServerClosePoolSize is the number of goroutines that will handle closing\n\t// server connections.\n\tServerClosePoolSize uint\n\n\t// ClientIdleTimeout is how long until we'll consider a client connection\n\t// idle and disconnect and release it's resources.\n\tClientIdleTimeout time.Duration\n\n\t// MaxPerClientConnections is how many client connections are allowed from a\n\t// single client.\n\tMaxPerClientConnections uint\n\n\t// GetLastErrorTimeout is how long we'll hold on to an acquired server\n\t// connection expecting a possibly getLastError call.\n\tGetLastErrorTimeout time.Duration\n\n\t// MessageTimeout is used to determine the timeout for a single message to be\n\t// proxied.\n\tMessageTimeout time.Duration\n\n\t// Name is the name of the replica set to connect to. Nodes that are not part\n\t// of this replica set will be ignored. If this is empty, the first replica set\n\t// will be used\n\tName string\n\n\tproxyToReal map[string]string\n\trealToProxy map[string]string\n\tignoredReal map[string]ReplicaState\n\tproxies     map[string]*Proxy\n\trestarter   *sync.Once\n\tlastState   *ReplicaSetState\n}\n\n// Start starts proxies to support this ReplicaSet.\nfunc (r *ReplicaSet) Start() error {\n\tr.proxyToReal = make(map[string]string)\n\tr.realToProxy = make(map[string]string)\n\tr.ignoredReal = make(map[string]ReplicaState)\n\tr.proxies = make(map[string]*Proxy)\n\n\tif r.Addrs == \"\" {\n\t\treturn errNoAddrsGiven\n\t}\n\n\trawAddrs := strings.Split(r.Addrs, \",\")\n\tvar err error\n\tr.lastState, err = r.ReplicaSetStateCreator.FromAddrs(rawAddrs, r.Name)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\thealthyAddrs := r.lastState.Addrs()\n\n\t// Ensure we have at least one health address.\n\tif len(healthyAddrs) == 0 {\n\t\treturn stackerr.Newf(\"no healthy primaries or secondaries: %s\", r.Addrs)\n\t}\n\n\t// Add discovered nodes to seed address list. Over time if the original seed\n\t// nodes have gone away and new nodes have joined this ensures that we'll\n\t// still be able to connect.\n\tr.Addrs = strings.Join(uniq(append(rawAddrs, healthyAddrs...)), \",\")\n\n\tr.restarter = new(sync.Once)\n\n\tfor _, addr := range healthyAddrs {\n\t\tlistener, err := r.newListener()\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tp := &Proxy{\n\t\t\tLog:            r.Log,\n\t\t\tReplicaSet:     r,\n\t\t\tClientListener: listener,\n\t\t\tProxyAddr:      r.proxyAddr(listener),\n\t\t\tMongoAddr:      addr,\n\t\t}\n\t\tif err := r.add(p); err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\t// add the ignored hosts, unless lastRS is nil (single node mode)\n\tif r.lastState.lastRS != nil {\n\t\tfor _, member := range r.lastState.lastRS.Members {\n\t\t\tif _, ok := r.realToProxy[member.Name]; !ok {\n\t\t\t\tr.ignoredReal[member.Name] = member.State\n\t\t\t}\n\t\t}\n\t}\n\n\tvar wg sync.WaitGroup\n\twg.Add(len(r.proxies))\n\terrch := make(chan error, len(r.proxies))\n\tfor _, p := range r.proxies {\n\t\tgo func(p *Proxy) {\n\t\t\tdefer wg.Done()\n\t\t\tif err := p.Start(); err != nil {\n\t\t\t\tr.Log.Error(err)\n\t\t\t\terrch <- stackerr.Wrap(err)\n\t\t\t}\n\t\t}(p)\n\t}\n\twg.Wait()\n\tselect {\n\tdefault:\n\t\treturn nil\n\tcase err := <-errch:\n\t\treturn err\n\t}\n}\n\n// Stop stops all the associated proxies for this ReplicaSet.\nfunc (r *ReplicaSet) Stop() error {\n\treturn r.stop(false)\n}\n\nfunc (r *ReplicaSet) stop(hard bool) error {\n\tvar wg sync.WaitGroup\n\twg.Add(len(r.proxies))\n\terrch := make(chan error, len(r.proxies))\n\tfor _, p := range r.proxies {\n\t\tgo func(p *Proxy) {\n\t\t\tdefer wg.Done()\n\t\t\tif err := p.stop(hard); err != nil {\n\t\t\t\tr.Log.Error(err)\n\t\t\t\terrch <- stackerr.Wrap(err)\n\t\t\t}\n\t\t}(p)\n\t}\n\twg.Wait()\n\tselect {\n\tdefault:\n\t\treturn nil\n\tcase err := <-errch:\n\t\treturn err\n\t}\n}\n\n// Restart stops all the proxies and restarts them. This is used when we detect\n// an RS config change, like when an election happens.\nfunc (r *ReplicaSet) Restart() {\n\tr.restarter.Do(func() {\n\t\tr.Log.Info(\"restart triggered\")\n\t\tif err := r.stop(*hardRestart); err != nil {\n\t\t\t// We log and ignore this hoping for a successful start anyways.\n\t\t\tr.Log.Errorf(\"stop failed for restart: %s\", err)\n\t\t} else {\n\t\t\tr.Log.Info(\"successfully stopped for restart\")\n\t\t}\n\n\t\tif err := r.Start(); err != nil {\n\t\t\t// We panic here because we can't repair from here and are pretty much\n\t\t\t// fucked.\n\t\t\tpanic(fmt.Errorf(\"start failed for restart: %s\", err))\n\t\t}\n\n\t\tr.Log.Info(\"successfully restarted\")\n\t})\n}\n\nfunc (r *ReplicaSet) proxyAddr(l net.Listener) string {\n\t_, port, err := net.SplitHostPort(l.Addr().String())\n\tif err != nil {\n\t\tpanic(err)\n\t}\n\n\treturn fmt.Sprintf(\"%s:%s\", r.proxyHostname(), port)\n}\n\nfunc (r *ReplicaSet) proxyHostname() string {\n\tconst home = \"127.0.0.1\"\n\n\thostname, err := os.Hostname()\n\tif err != nil {\n\t\tr.Log.Error(err)\n\t\treturn home\n\t}\n\n\t// The follow logic ensures that the hostname resolves to a local address.\n\t// If it doesn't we don't use it since it probably wont work anyways.\n\thostnameAddrs, err := net.LookupHost(hostname)\n\tif err != nil {\n\t\tr.Log.Error(err)\n\t\treturn home\n\t}\n\n\tinterfaceAddrs, err := net.InterfaceAddrs()\n\tif err != nil {\n\t\tr.Log.Error(err)\n\t\treturn home\n\t}\n\n\tfor _, ia := range interfaceAddrs {\n\t\tsa := ia.String()\n\t\tfor _, ha := range hostnameAddrs {\n\t\t\t// check for an exact match or a match ignoring the suffix bits\n\t\t\tif sa == ha || strings.HasPrefix(sa, ha+\"/\") {\n\t\t\t\treturn hostname\n\t\t\t}\n\t\t}\n\t}\n\tr.Log.Warnf(\"hostname %s doesn't resolve to the current host\", hostname)\n\treturn home\n}\n\nfunc (r *ReplicaSet) newListener() (net.Listener, error) {\n\tfor i := r.PortStart; i <= r.PortEnd; i++ {\n\t\tlistener, err := net.Listen(\"tcp\", fmt.Sprintf(\":%d\", i))\n\t\tif err == nil {\n\t\t\treturn listener, nil\n\t\t}\n\t}\n\treturn nil, fmt.Errorf(\n\t\t\"could not find a free port in range %d-%d\",\n\t\tr.PortStart,\n\t\tr.PortEnd,\n\t)\n}\n\n// add a proxy/mongo mapping.\nfunc (r *ReplicaSet) add(p *Proxy) error {\n\tif _, ok := r.proxyToReal[p.ProxyAddr]; ok {\n\t\treturn fmt.Errorf(\"proxy %s already used in ReplicaSet\", p.ProxyAddr)\n\t}\n\tif _, ok := r.realToProxy[p.MongoAddr]; ok {\n\t\treturn fmt.Errorf(\"mongo %s already exists in ReplicaSet\", p.MongoAddr)\n\t}\n\tr.Log.Infof(\"added %s\", p)\n\tr.proxyToReal[p.ProxyAddr] = p.MongoAddr\n\tr.realToProxy[p.MongoAddr] = p.ProxyAddr\n\tr.proxies[p.ProxyAddr] = p\n\treturn nil\n}\n\n// Proxy returns the corresponding proxy address for the given real mongo\n// address.\nfunc (r *ReplicaSet) Proxy(h string) (string, error) {\n\tp, ok := r.realToProxy[h]\n\tif !ok {\n\t\tif s, ok := r.ignoredReal[h]; ok {\n\t\t\treturn \"\", &ProxyMapperError{\n\t\t\t\tRealHost: h,\n\t\t\t\tState:    s,\n\t\t\t}\n\t\t}\n\t\treturn \"\", fmt.Errorf(\"mongo %s is not in ReplicaSet\", h)\n\t}\n\treturn p, nil\n}\n\n// ProxyMembers returns the list of proxy members in this ReplicaSet.\nfunc (r *ReplicaSet) ProxyMembers() []string {\n\tmembers := make([]string, 0, len(r.proxyToReal))\n\tfor r := range r.proxyToReal {\n\t\tmembers = append(members, r)\n\t}\n\treturn members\n}\n\n// SameRS checks if the given replSetGetStatusResponse is the same as the last\n// state.\nfunc (r *ReplicaSet) SameRS(o *replSetGetStatusResponse) bool {\n\treturn r.lastState.SameRS(o)\n}\n\n// SameIM checks if the given isMasterResponse is the same as the last state.\nfunc (r *ReplicaSet) SameIM(o *isMasterResponse) bool {\n\treturn r.lastState.SameIM(o)\n}\n\n// ProxyMapperError occurs when a known host is being ignored and does not have\n// a corresponding proxy address.\ntype ProxyMapperError struct {\n\tRealHost string\n\tState    ReplicaState\n}\n\nfunc (p *ProxyMapperError) Error() string {\n\treturn fmt.Sprintf(\"error mapping host %s in state %s\", p.RealHost, p.State)\n}\n\n// uniq takes a slice of strings and returns a new slice with duplicates\n// removed.\nfunc uniq(set []string) []string {\n\tm := make(map[string]struct{}, len(set))\n\tfor _, s := range set {\n\t\tm[s] = struct{}{}\n\t}\n\tnews := make([]string, 0, len(m))\n\tfor s := range m {\n\t\tnews = append(news, s)\n\t}\n\treturn news\n}\n"
  },
  {
    "path": "replica_set_test.go",
    "content": "package dvara\n\nimport (\n\t\"fmt\"\n\t\"testing\"\n\n\t\"github.com/facebookgo/subset\"\n\n\t\"gopkg.in/mgo.v2\"\n\t\"gopkg.in/mgo.v2/bson\"\n)\n\nfunc TestReplicaSetMembers(t *testing.T) {\n\tt.Parallel()\n\th := NewReplicaSetHarness(3, t)\n\tdefer h.Stop()\n\n\tproxyMembers := h.ReplicaSet.ProxyMembers()\n\tsession := h.ProxySession()\n\tdefer session.Close()\n\tstatus, err := replSetGetStatus(session)\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\nouterProxyResponseCheckLoop:\n\tfor _, m := range status.Members {\n\t\tfor _, p := range proxyMembers {\n\t\t\tif m.Name == p {\n\t\t\t\tcontinue outerProxyResponseCheckLoop\n\t\t\t}\n\t\t}\n\t\tt.Fatalf(\"Unexpected member: %s\", m.Name)\n\t}\n}\n\nfunc TestStopNodeInReplica(t *testing.T) {\n\tt.Parallel()\n\th := NewReplicaSetHarness(2, t)\n\tdefer h.Stop()\n\n\tconst dbName = \"test\"\n\tconst colName = \"foo\"\n\tconst keyName = \"answer\"\n\td := bson.M{\"answer\": \"42\"}\n\ts := h.ProxySession()\n\tdefer s.Close()\n\ts.SetSafe(&mgo.Safe{W: 2, WMode: \"majority\"})\n\tif err := s.DB(dbName).C(colName).Insert(d); err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\th.MgoReplicaSet.Servers[0].Stop()\n\n\ts.SetMode(mgo.Monotonic, true)\n\tvar actual bson.M\n\tif err := s.DB(dbName).C(colName).Find(d).One(&actual); err != nil {\n\t\tt.Fatal(err)\n\t}\n\n\tsubset.Assert(t, d, actual)\n}\n\nfunc TestProxyNotInReplicaSet(t *testing.T) {\n\tt.Parallel()\n\th := NewSingleHarness(t)\n\tdefer h.Stop()\n\taddr := \"127.0.0.1:666\"\n\texpected := fmt.Sprintf(\"mongo %s is not in ReplicaSet\", addr)\n\t_, err := h.ReplicaSet.Proxy(addr)\n\tif err == nil || err.Error() != expected {\n\t\tt.Fatalf(\"did not get expected error, got: %s\", err)\n\t}\n}\n\nfunc TestAddSameProxyToReplicaSet(t *testing.T) {\n\tt.Parallel()\n\tr := &ReplicaSet{\n\t\tLog:         &tLogger{TB: t},\n\t\tproxyToReal: make(map[string]string),\n\t\trealToProxy: make(map[string]string),\n\t\tproxies:     make(map[string]*Proxy),\n\t}\n\tp := &Proxy{\n\t\tProxyAddr: \"1\",\n\t\tMongoAddr: \"2\",\n\t}\n\tif err := r.add(p); err != nil {\n\t\tt.Fatal(err)\n\t}\n\texpected := fmt.Sprintf(\"proxy %s already used in ReplicaSet\", p.ProxyAddr)\n\terr := r.add(p)\n\tif err == nil || err.Error() != expected {\n\t\tt.Fatalf(\"did not get expected error, got: %s\", err)\n\t}\n}\n\nfunc TestAddSameMongoToReplicaSet(t *testing.T) {\n\tt.Parallel()\n\tr := &ReplicaSet{\n\t\tLog:         &tLogger{TB: t},\n\t\tproxyToReal: make(map[string]string),\n\t\trealToProxy: make(map[string]string),\n\t\tproxies:     make(map[string]*Proxy),\n\t}\n\tp := &Proxy{\n\t\tProxyAddr: \"1\",\n\t\tMongoAddr: \"2\",\n\t}\n\tif err := r.add(p); err != nil {\n\t\tt.Fatal(err)\n\t}\n\tp = &Proxy{\n\t\tProxyAddr: \"3\",\n\t\tMongoAddr: p.MongoAddr,\n\t}\n\texpected := fmt.Sprintf(\"mongo %s already exists in ReplicaSet\", p.MongoAddr)\n\terr := r.add(p)\n\tif err == nil || err.Error() != expected {\n\t\tt.Fatalf(\"did not get expected error, got: %s\", err)\n\t}\n}\n\nfunc TestNewListenerZeroZeroRandomPort(t *testing.T) {\n\tt.Parallel()\n\tr := &ReplicaSet{}\n\tl, err := r.newListener()\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tl.Close()\n}\n\nfunc TestNewListenerError(t *testing.T) {\n\tt.Parallel()\n\tr := &ReplicaSet{PortStart: 1, PortEnd: 1}\n\t_, err := r.newListener()\n\texpected := \"could not find a free port in range 1-1\"\n\tif err == nil || err.Error() != expected {\n\t\tt.Fatalf(\"did not get expected error, got: %s\", err)\n\t}\n}\n"
  },
  {
    "path": "response_rewriter.go",
    "content": "package dvara\n\nimport (\n\t\"bytes\"\n\t\"errors\"\n\t\"flag\"\n\t\"fmt\"\n\t\"io\"\n\t\"io/ioutil\"\n\t\"strings\"\n\n\t\"github.com/davecgh/go-spew/spew\"\n\n\t\"gopkg.in/mgo.v2/bson\"\n)\n\nvar (\n\tproxyAllQueries = flag.Bool(\n\t\t\"dvara.proxy-all\",\n\t\tfalse,\n\t\t\"if true all queries will be proxied and logger\",\n\t)\n\n\tadminCollectionName = []byte(\"admin.$cmd\\000\")\n\tcmdCollectionSuffix = []byte(\".$cmd\\000\")\n)\n\n// ProxyQuery proxies an OpQuery and a corresponding response.\ntype ProxyQuery struct {\n\tLog                              Logger                            `inject:\"\"`\n\tGetLastErrorRewriter             *GetLastErrorRewriter             `inject:\"\"`\n\tIsMasterResponseRewriter         *IsMasterResponseRewriter         `inject:\"\"`\n\tReplSetGetStatusResponseRewriter *ReplSetGetStatusResponseRewriter `inject:\"\"`\n}\n\n// Proxy proxies an OpQuery and a corresponding response.\nfunc (p *ProxyQuery) Proxy(\n\th *messageHeader,\n\tclient io.ReadWriter,\n\tserver io.ReadWriter,\n\tlastError *LastError,\n) error {\n\n\t// https://github.com/mongodb/mongo/search?q=lastError.disableForCommand\n\t// Shows the logic we need to be in sync with. Unfortunately it isn't a\n\t// simple check to determine this, and may change underneath us at the mongo\n\t// layer.\n\tresetLastError := true\n\n\tparts := [][]byte{h.ToWire()}\n\n\tvar flags [4]byte\n\tif _, err := io.ReadFull(client, flags[:]); err != nil {\n\t\tp.Log.Error(err)\n\t\treturn err\n\t}\n\tparts = append(parts, flags[:])\n\n\tfullCollectionName, err := readCString(client)\n\tif err != nil {\n\t\tp.Log.Error(err)\n\t\treturn err\n\t}\n\tparts = append(parts, fullCollectionName)\n\n\tvar rewriter responseRewriter\n\tif *proxyAllQueries || bytes.HasSuffix(fullCollectionName, cmdCollectionSuffix) {\n\t\tvar twoInt32 [8]byte\n\t\tif _, err := io.ReadFull(client, twoInt32[:]); err != nil {\n\t\t\tp.Log.Error(err)\n\t\t\treturn err\n\t\t}\n\t\tparts = append(parts, twoInt32[:])\n\n\t\tqueryDoc, err := readDocument(client)\n\t\tif err != nil {\n\t\t\tp.Log.Error(err)\n\t\t\treturn err\n\t\t}\n\t\tparts = append(parts, queryDoc)\n\n\t\tvar q bson.D\n\t\tif err := bson.Unmarshal(queryDoc, &q); err != nil {\n\t\t\tp.Log.Error(err)\n\t\t\treturn err\n\t\t}\n\n\t\tp.Log.Debugf(\n\t\t\t\"buffered OpQuery for %s: %s\",\n\t\t\tfullCollectionName[:len(fullCollectionName)-1],\n\t\t\tspew.Sdump(q),\n\t\t)\n\n\t\tif hasKey(q, \"getLastError\") {\n\t\t\treturn p.GetLastErrorRewriter.Rewrite(\n\t\t\t\th,\n\t\t\t\tparts,\n\t\t\t\tclient,\n\t\t\t\tserver,\n\t\t\t\tlastError,\n\t\t\t)\n\t\t}\n\n\t\tif hasKey(q, \"isMaster\") {\n\t\t\trewriter = p.IsMasterResponseRewriter\n\t\t}\n\t\tif bytes.Equal(adminCollectionName, fullCollectionName) && hasKey(q, \"replSetGetStatus\") {\n\t\t\trewriter = p.ReplSetGetStatusResponseRewriter\n\t\t}\n\n\t\tif rewriter != nil {\n\t\t\t// If forShell is specified, we don't want to reset the last error. See\n\t\t\t// comment above around resetLastError for details.\n\t\t\tresetLastError = hasKey(q, \"forShell\")\n\t\t}\n\t}\n\n\tif resetLastError && lastError.Exists() {\n\t\tp.Log.Debug(\"reset getLastError cache\")\n\t\tlastError.Reset()\n\t}\n\n\tvar written int\n\tfor _, b := range parts {\n\t\tn, err := server.Write(b)\n\t\tif err != nil {\n\t\t\tp.Log.Error(err)\n\t\t\treturn err\n\t\t}\n\t\twritten += n\n\t}\n\n\tpending := int64(h.MessageLength) - int64(written)\n\tif _, err := io.CopyN(server, client, pending); err != nil {\n\t\tp.Log.Error(err)\n\t\treturn err\n\t}\n\n\tif rewriter != nil {\n\t\tif err := rewriter.Rewrite(client, server); err != nil {\n\t\t\treturn err\n\t\t}\n\t\treturn nil\n\t}\n\n\tif err := copyMessage(client, server); err != nil {\n\t\tp.Log.Error(err)\n\t\treturn err\n\t}\n\n\treturn nil\n}\n\n// LastError holds the last known error.\ntype LastError struct {\n\theader *messageHeader\n\trest   bytes.Buffer\n}\n\n// Exists returns true if this instance contains a cached error.\nfunc (l *LastError) Exists() bool {\n\treturn l.header != nil\n}\n\n// Reset resets the stored error clearing it.\nfunc (l *LastError) Reset() {\n\tl.header = nil\n\tl.rest.Reset()\n}\n\n// GetLastErrorRewriter handles getLastError requests and proxies, caches or\n// sends cached responses as necessary.\ntype GetLastErrorRewriter struct {\n\tLog Logger `inject:\"\"`\n}\n\n// Rewrite handles getLastError requests.\nfunc (r *GetLastErrorRewriter) Rewrite(\n\th *messageHeader,\n\tparts [][]byte,\n\tclient io.ReadWriter,\n\tserver io.ReadWriter,\n\tlastError *LastError,\n) error {\n\n\tif !lastError.Exists() {\n\t\t// We're going to be performing a real getLastError query and caching the\n\t\t// response.\n\t\tvar written int\n\t\tfor _, b := range parts {\n\t\t\tn, err := server.Write(b)\n\t\t\tif err != nil {\n\t\t\t\tr.Log.Error(err)\n\t\t\t\treturn err\n\t\t\t}\n\t\t\twritten += n\n\t\t}\n\n\t\tpending := int64(h.MessageLength) - int64(written)\n\t\tif _, err := io.CopyN(server, client, pending); err != nil {\n\t\t\tr.Log.Error(err)\n\t\t\treturn err\n\t\t}\n\n\t\tvar err error\n\t\tif lastError.header, err = readHeader(server); err != nil {\n\t\t\tr.Log.Error(err)\n\t\t\treturn err\n\t\t}\n\t\tpending = int64(lastError.header.MessageLength - headerLen)\n\t\tif _, err = io.CopyN(&lastError.rest, server, pending); err != nil {\n\t\t\tr.Log.Error(err)\n\t\t\treturn err\n\t\t}\n\t\tr.Log.Debugf(\"caching new getLastError response: %s\", lastError.rest.Bytes())\n\t} else {\n\t\t// We need to discard the pending bytes from the client from the query\n\t\t// before we send it our cached response.\n\t\tvar written int\n\t\tfor _, b := range parts {\n\t\t\twritten += len(b)\n\t\t}\n\t\tpending := int64(h.MessageLength) - int64(written)\n\t\tif _, err := io.CopyN(ioutil.Discard, client, pending); err != nil {\n\t\t\tr.Log.Error(err)\n\t\t\treturn err\n\t\t}\n\t\t// Modify and send the cached response for this request.\n\t\tlastError.header.ResponseTo = h.RequestID\n\t\tr.Log.Debugf(\"using cached getLastError response: %s\", lastError.rest.Bytes())\n\t}\n\n\tif err := lastError.header.WriteTo(client); err != nil {\n\t\tr.Log.Error(err)\n\t\treturn err\n\t}\n\tif _, err := client.Write(lastError.rest.Bytes()); err != nil {\n\t\tr.Log.Error(err)\n\t\treturn err\n\t}\n\n\treturn nil\n}\n\nvar errRSChanged = errors.New(\"dvara: replset config changed\")\n\n// ProxyMapper maps real mongo addresses to their corresponding proxy\n// addresses.\ntype ProxyMapper interface {\n\tProxy(h string) (string, error)\n}\n\n// ReplicaStateCompare provides the last ReplicaSetState and allows for\n// checking if it has changed as we rewrite/proxy the isMaster &\n// replSetGetStatus queries.\ntype ReplicaStateCompare interface {\n\tSameRS(o *replSetGetStatusResponse) bool\n\tSameIM(o *isMasterResponse) bool\n}\n\ntype responseRewriter interface {\n\tRewrite(client io.Writer, server io.Reader) error\n}\n\ntype replyPrefix [20]byte\n\nvar emptyPrefix replyPrefix\n\n// ReplyRW provides common helpers for rewriting replies from the server.\ntype ReplyRW struct {\n\tLog Logger `inject:\"\"`\n}\n\n// ReadOne reads a 1 document response, from the server, unmarshals it into v\n// and returns the various parts.\nfunc (r *ReplyRW) ReadOne(server io.Reader, v interface{}) (*messageHeader, replyPrefix, int32, error) {\n\th, err := readHeader(server)\n\tif err != nil {\n\t\tr.Log.Error(err)\n\t\treturn nil, emptyPrefix, 0, err\n\t}\n\n\tif h.OpCode != OpReply {\n\t\terr := fmt.Errorf(\"readOneReplyDoc: expected op %s, got %s\", OpReply, h.OpCode)\n\t\treturn nil, emptyPrefix, 0, err\n\t}\n\n\tvar prefix replyPrefix\n\tif _, err := io.ReadFull(server, prefix[:]); err != nil {\n\t\tr.Log.Error(err)\n\t\treturn nil, emptyPrefix, 0, err\n\t}\n\n\tnumDocs := getInt32(prefix[:], 16)\n\tif numDocs != 1 {\n\t\terr := fmt.Errorf(\"readOneReplyDoc: can only handle 1 result document, got: %d\", numDocs)\n\t\treturn nil, emptyPrefix, 0, err\n\t}\n\n\trawDoc, err := readDocument(server)\n\tif err != nil {\n\t\tr.Log.Error(err)\n\t\treturn nil, emptyPrefix, 0, err\n\t}\n\n\tif err := bson.Unmarshal(rawDoc, v); err != nil {\n\t\tr.Log.Error(err)\n\t\treturn nil, emptyPrefix, 0, err\n\t}\n\n\treturn h, prefix, int32(len(rawDoc)), nil\n}\n\n// WriteOne writes a rewritten response to the client.\nfunc (r *ReplyRW) WriteOne(client io.Writer, h *messageHeader, prefix replyPrefix, oldDocLen int32, v interface{}) error {\n\tnewDoc, err := bson.Marshal(v)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\th.MessageLength = h.MessageLength - oldDocLen + int32(len(newDoc))\n\tparts := [][]byte{h.ToWire(), prefix[:], newDoc}\n\tfor _, p := range parts {\n\t\tif _, err := client.Write(p); err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\treturn nil\n}\n\ntype isMasterResponse struct {\n\tHosts   []string `bson:\"hosts,omitempty\"`\n\tPrimary string   `bson:\"primary,omitempty\"`\n\tMe      string   `bson:\"me,omitempty\"`\n\tExtra   bson.M   `bson:\",inline\"`\n}\n\n// IsMasterResponseRewriter rewrites the response for the \"isMaster\" query.\ntype IsMasterResponseRewriter struct {\n\tLog                 Logger              `inject:\"\"`\n\tProxyMapper         ProxyMapper         `inject:\"\"`\n\tReplyRW             *ReplyRW            `inject:\"\"`\n\tReplicaStateCompare ReplicaStateCompare `inject:\"\"`\n}\n\n// Rewrite rewrites the response for the \"isMaster\" query.\nfunc (r *IsMasterResponseRewriter) Rewrite(client io.Writer, server io.Reader) error {\n\tvar err error\n\tvar q isMasterResponse\n\th, prefix, docLen, err := r.ReplyRW.ReadOne(server, &q)\n\tif err != nil {\n\t\treturn err\n\t}\n\tif !r.ReplicaStateCompare.SameIM(&q) {\n\t\treturn errRSChanged\n\t}\n\n\tvar newHosts []string\n\tfor _, h := range q.Hosts {\n\t\tnewH, err := r.ProxyMapper.Proxy(h)\n\t\tif err != nil {\n\t\t\tif pme, ok := err.(*ProxyMapperError); ok {\n\t\t\t\tif pme.State != ReplicaStateArbiter {\n\t\t\t\t\tr.Log.Errorf(\"dropping member %s in state %s\", h, pme.State)\n\t\t\t\t}\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\t// unknown err\n\t\t\treturn err\n\t\t}\n\t\tnewHosts = append(newHosts, newH)\n\t}\n\tq.Hosts = newHosts\n\n\tif q.Primary != \"\" {\n\t\t// failure in mapping the primary is fatal\n\t\tif q.Primary, err = r.ProxyMapper.Proxy(q.Primary); err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\tif q.Me != \"\" {\n\t\t// failure in mapping me is fatal\n\t\tif q.Me, err = r.ProxyMapper.Proxy(q.Me); err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\treturn r.ReplyRW.WriteOne(client, h, prefix, docLen, q)\n}\n\ntype statusMember struct {\n\tName  string       `bson:\"name\"`\n\tState ReplicaState `bson:\"stateStr,omitempty\"`\n\tSelf  bool         `bson:\"self,omitempty\"`\n\tExtra bson.M       `bson:\",inline\"`\n}\n\ntype replSetGetStatusResponse struct {\n\tName    string                 `bson:\"set,omitempty\"`\n\tMembers []statusMember         `bson:\"members\"`\n\tExtra   map[string]interface{} `bson:\",inline\"`\n}\n\n// ReplSetGetStatusResponseRewriter rewrites the \"replSetGetStatus\" response.\ntype ReplSetGetStatusResponseRewriter struct {\n\tLog                 Logger              `inject:\"\"`\n\tProxyMapper         ProxyMapper         `inject:\"\"`\n\tReplyRW             *ReplyRW            `inject:\"\"`\n\tReplicaStateCompare ReplicaStateCompare `inject:\"\"`\n}\n\n// Rewrite rewrites the \"replSetGetStatus\" response.\nfunc (r *ReplSetGetStatusResponseRewriter) Rewrite(client io.Writer, server io.Reader) error {\n\tvar err error\n\tvar q replSetGetStatusResponse\n\th, prefix, docLen, err := r.ReplyRW.ReadOne(server, &q)\n\tif err != nil {\n\t\treturn err\n\t}\n\tif !r.ReplicaStateCompare.SameRS(&q) {\n\t\treturn errRSChanged\n\t}\n\n\tvar newMembers []statusMember\n\tfor _, m := range q.Members {\n\t\tnewH, err := r.ProxyMapper.Proxy(m.Name)\n\t\tif err != nil {\n\t\t\tif pme, ok := err.(*ProxyMapperError); ok {\n\t\t\t\tif pme.State != ReplicaStateArbiter {\n\t\t\t\t\tr.Log.Errorf(\"dropping member %s in state %s\", h, pme.State)\n\t\t\t\t}\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\t// unknown err\n\t\t\treturn err\n\t\t}\n\t\tm.Name = newH\n\t\tnewMembers = append(newMembers, m)\n\t}\n\tq.Members = newMembers\n\treturn r.ReplyRW.WriteOne(client, h, prefix, docLen, q)\n}\n\n// case insensitive check for the specified key name in the top level.\nfunc hasKey(d bson.D, k string) bool {\n\tfor _, v := range d {\n\t\tif strings.EqualFold(v.Name, k) {\n\t\t\treturn true\n\t\t}\n\t}\n\treturn false\n}\n"
  },
  {
    "path": "response_rewriter_test.go",
    "content": "package dvara\n\nimport (\n\t\"bytes\"\n\t\"errors\"\n\t\"io\"\n\t\"reflect\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/davecgh/go-spew/spew\"\n\t\"github.com/facebookgo/ensure\"\n\t\"github.com/facebookgo/inject\"\n\t\"github.com/facebookgo/startstop\"\n\n\t\"gopkg.in/mgo.v2/bson\"\n)\n\nvar errInvalidBSON = errors.New(\"invalid BSON\")\n\ntype invalidBSON int\n\nfunc (i invalidBSON) GetBSON() (interface{}, error) {\n\treturn nil, errInvalidBSON\n}\n\nvar errProxyNotFound = errors.New(\"proxy not found\")\n\ntype fakeProxyMapper struct {\n\tm map[string]string\n}\n\nfunc (t fakeProxyMapper) Proxy(h string) (string, error) {\n\tif t.m != nil {\n\t\tif r, ok := t.m[h]; ok {\n\t\t\treturn r, nil\n\t\t}\n\t}\n\treturn \"\", errProxyNotFound\n}\n\ntype fakeReplicaStateCompare struct{ sameRS, sameIM bool }\n\nfunc (f fakeReplicaStateCompare) SameRS(o *replSetGetStatusResponse) bool {\n\treturn f.sameRS\n}\n\nfunc (f fakeReplicaStateCompare) SameIM(o *isMasterResponse) bool {\n\treturn f.sameIM\n}\n\nfunc fakeReader(h messageHeader, rest []byte) io.Reader {\n\treturn bytes.NewReader(append(h.ToWire(), rest...))\n}\n\nfunc fakeSingleDocReply(v interface{}) io.Reader {\n\tb, err := bson.Marshal(v)\n\tif err != nil {\n\t\tpanic(err)\n\t}\n\tb = append(\n\t\t[]byte{\n\t\t\t0, 0, 0, 0,\n\t\t\t0, 0, 0, 0, 0, 0, 0, 0,\n\t\t\t0, 0, 0, 0,\n\t\t\t1, 0, 0, 0,\n\t\t},\n\t\tb...,\n\t)\n\th := messageHeader{\n\t\tOpCode:        OpReply,\n\t\tMessageLength: int32(headerLen + len(b)),\n\t}\n\treturn fakeReader(h, b)\n}\n\ntype fakeReadWriter struct {\n\tio.Reader\n\tio.Writer\n}\n\nfunc TestResponseRWReadOne(t *testing.T) {\n\tt.Parallel()\n\tcases := []struct {\n\t\tName   string\n\t\tServer io.Reader\n\t\tError  string\n\t}{\n\t\t{\n\t\t\tName:   \"no header\",\n\t\t\tServer: bytes.NewReader(nil),\n\t\t\tError:  \"EOF\",\n\t\t},\n\t\t{\n\t\t\tName:   \"non reply op\",\n\t\t\tServer: bytes.NewReader((messageHeader{OpCode: OpDelete}).ToWire()),\n\t\t\tError:  \"expected op REPLY, got DELETE\",\n\t\t},\n\t\t{\n\t\t\tName:   \"EOF before flags\",\n\t\t\tServer: bytes.NewReader((messageHeader{OpCode: OpReply}).ToWire()),\n\t\t\tError:  \"EOF\",\n\t\t},\n\t\t{\n\t\t\tName: \"more than 1 document\",\n\t\t\tServer: fakeReader(\n\t\t\t\tmessageHeader{OpCode: OpReply},\n\t\t\t\t[]byte{\n\t\t\t\t\t0, 0, 0, 0,\n\t\t\t\t\t0, 0, 0, 0, 0, 0, 0, 0,\n\t\t\t\t\t0, 0, 0, 0,\n\t\t\t\t\t2, 0, 0, 0,\n\t\t\t\t},\n\t\t\t),\n\t\t\tError: \"can only handle 1 result document, got: 2\",\n\t\t},\n\t\t{\n\t\t\tName: \"EOF before document\",\n\t\t\tServer: fakeReader(\n\t\t\t\tmessageHeader{OpCode: OpReply},\n\t\t\t\t[]byte{\n\t\t\t\t\t0, 0, 0, 0,\n\t\t\t\t\t0, 0, 0, 0, 0, 0, 0, 0,\n\t\t\t\t\t0, 0, 0, 0,\n\t\t\t\t\t1, 0, 0, 0,\n\t\t\t\t},\n\t\t\t),\n\t\t\tError: \"EOF\",\n\t\t},\n\t\t{\n\t\t\tName: \"corrupted document\",\n\t\t\tServer: fakeReader(\n\t\t\t\tmessageHeader{OpCode: OpReply},\n\t\t\t\t[]byte{\n\t\t\t\t\t0, 0, 0, 0,\n\t\t\t\t\t0, 0, 0, 0, 0, 0, 0, 0,\n\t\t\t\t\t0, 0, 0, 0,\n\t\t\t\t\t1, 0, 0, 0,\n\t\t\t\t\t5, 0, 0, 0,\n\t\t\t\t\t1,\n\t\t\t\t},\n\t\t\t),\n\t\t\tError: \"Document is corrupted\",\n\t\t},\n\t}\n\n\tfor _, c := range cases {\n\t\tr := &ReplyRW{Log: &tLogger{TB: t}}\n\t\tm := bson.M{}\n\t\t_, _, _, err := r.ReadOne(c.Server, m)\n\t\tif err == nil {\n\t\t\tt.Errorf(\"was expecting an error for case %s\", c.Name)\n\t\t}\n\t\tif !strings.Contains(err.Error(), c.Error) {\n\t\t\tt.Errorf(\"did not get expected error for case %s instead got %s\", c.Name, err)\n\t\t}\n\t}\n}\n\nfunc TestResponseRWWriteOne(t *testing.T) {\n\terrWrite := errors.New(\"write error\")\n\tt.Parallel()\n\tcases := []struct {\n\t\tName   string\n\t\tClient io.Writer\n\t\tHeader messageHeader\n\t\tPrefix replyPrefix\n\t\tDocLen int32\n\t\tValue  interface{}\n\t\tError  string\n\t}{\n\t\t{\n\t\t\tName:  \"invalid bson\",\n\t\t\tValue: invalidBSON(0),\n\t\t\tError: errInvalidBSON.Error(),\n\t\t},\n\t\t{\n\t\t\tName:  \"write error\",\n\t\t\tValue: map[string]string{},\n\t\t\tClient: testWriter{\n\t\t\t\twrite: func(b []byte) (int, error) {\n\t\t\t\t\treturn 0, errWrite\n\t\t\t\t},\n\t\t\t},\n\t\t\tError: errWrite.Error(),\n\t\t},\n\t}\n\n\tfor _, c := range cases {\n\t\tr := &ReplyRW{Log: &tLogger{TB: t}}\n\t\terr := r.WriteOne(c.Client, &c.Header, c.Prefix, c.DocLen, c.Value)\n\t\tif err == nil {\n\t\t\tt.Errorf(\"was expecting an error for case %s\", c.Name)\n\t\t}\n\t\tif !strings.Contains(err.Error(), c.Error) {\n\t\t\tt.Errorf(\"did not get expected error for case %s instead got %s\", c.Name, err)\n\t\t}\n\t}\n}\n\nfunc TestIsMasterResponseRewriterFailures(t *testing.T) {\n\tt.Parallel()\n\tcases := []struct {\n\t\tName                string\n\t\tClient              io.Writer\n\t\tServer              io.Reader\n\t\tProxyMapper         ProxyMapper\n\t\tReplicaStateCompare ReplicaStateCompare\n\t\tError               string\n\t}{\n\t\t{\n\t\t\tName:   \"no header\",\n\t\t\tServer: bytes.NewReader(nil),\n\t\t\tError:  \"EOF\",\n\t\t},\n\t\t{\n\t\t\tName: \"unknown host in 'hosts'\",\n\t\t\tServer: fakeSingleDocReply(\n\t\t\t\tmap[string]interface{}{\n\t\t\t\t\t\"hosts\": []string{\"foo\"},\n\t\t\t\t},\n\t\t\t),\n\t\t\tError:               errProxyNotFound.Error(),\n\t\t\tProxyMapper:         fakeProxyMapper{},\n\t\t\tReplicaStateCompare: fakeReplicaStateCompare{sameIM: true, sameRS: true},\n\t\t},\n\t\t{\n\t\t\tName: \"unknown host in 'primary'\",\n\t\t\tServer: fakeSingleDocReply(\n\t\t\t\tmap[string]interface{}{\n\t\t\t\t\t\"primary\": \"foo\",\n\t\t\t\t},\n\t\t\t),\n\t\t\tError:               errProxyNotFound.Error(),\n\t\t\tProxyMapper:         fakeProxyMapper{},\n\t\t\tReplicaStateCompare: fakeReplicaStateCompare{sameIM: true, sameRS: true},\n\t\t},\n\t\t{\n\t\t\tName: \"unknown host in 'me'\",\n\t\t\tServer: fakeSingleDocReply(\n\t\t\t\tmap[string]interface{}{\n\t\t\t\t\t\"me\": \"foo\",\n\t\t\t\t},\n\t\t\t),\n\t\t\tError:               errProxyNotFound.Error(),\n\t\t\tProxyMapper:         fakeProxyMapper{},\n\t\t\tReplicaStateCompare: fakeReplicaStateCompare{sameIM: true, sameRS: true},\n\t\t},\n\t\t{\n\t\t\tName:                \"different im\",\n\t\t\tServer:              fakeSingleDocReply(map[string]interface{}{}),\n\t\t\tError:               errRSChanged.Error(),\n\t\t\tProxyMapper:         nil,\n\t\t\tReplicaStateCompare: fakeReplicaStateCompare{sameIM: false, sameRS: true},\n\t\t},\n\t}\n\n\tfor _, c := range cases {\n\t\tr := &IsMasterResponseRewriter{\n\t\t\tLog:                 &tLogger{TB: t},\n\t\t\tProxyMapper:         c.ProxyMapper,\n\t\t\tReplicaStateCompare: c.ReplicaStateCompare,\n\t\t\tReplyRW: &ReplyRW{\n\t\t\t\tLog: &tLogger{TB: t},\n\t\t\t},\n\t\t}\n\t\terr := r.Rewrite(c.Client, c.Server)\n\t\tif err == nil {\n\t\t\tt.Errorf(\"was expecting an error for case %s\", c.Name)\n\t\t}\n\t\tif !strings.Contains(err.Error(), c.Error) {\n\t\t\tt.Errorf(\"did not get expected error for case %s instead got %s\", c.Name, err)\n\t\t}\n\t}\n}\n\nfunc TestIsMasterResponseRewriterSuccess(t *testing.T) {\n\tproxyMapper := fakeProxyMapper{\n\t\tm: map[string]string{\n\t\t\t\"a\": \"1\",\n\t\t\t\"b\": \"2\",\n\t\t\t\"c\": \"3\",\n\t\t},\n\t}\n\tin := bson.M{\n\t\t\"hosts\":   []interface{}{\"a\", \"b\", \"c\"},\n\t\t\"me\":      \"a\",\n\t\t\"primary\": \"b\",\n\t\t\"foo\":     \"bar\",\n\t}\n\tout := bson.M{\n\t\t\"hosts\":   []interface{}{\"1\", \"2\", \"3\"},\n\t\t\"me\":      \"1\",\n\t\t\"primary\": \"2\",\n\t\t\"foo\":     \"bar\",\n\t}\n\tr := &IsMasterResponseRewriter{\n\t\tLog:                 &tLogger{TB: t},\n\t\tProxyMapper:         proxyMapper,\n\t\tReplicaStateCompare: fakeReplicaStateCompare{sameIM: true, sameRS: true},\n\t\tReplyRW: &ReplyRW{\n\t\t\tLog: &tLogger{TB: t},\n\t\t},\n\t}\n\n\tvar client bytes.Buffer\n\tif err := r.Rewrite(&client, fakeSingleDocReply(in)); err != nil {\n\t\tt.Fatal(err)\n\t}\n\tactualOut := bson.M{}\n\tdoc := client.Bytes()[headerLen+len(emptyPrefix):]\n\tif err := bson.Unmarshal(doc, &actualOut); err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif !reflect.DeepEqual(out, actualOut) {\n\t\tspew.Dump(out)\n\t\tspew.Dump(actualOut)\n\t\tt.Fatal(\"did not get expected output\")\n\t}\n}\n\nfunc TestReplSetGetStatusResponseRewriterFailures(t *testing.T) {\n\tt.Parallel()\n\tcases := []struct {\n\t\tName                string\n\t\tClient              io.Writer\n\t\tServer              io.Reader\n\t\tProxyMapper         ProxyMapper\n\t\tReplicaStateCompare ReplicaStateCompare\n\t\tError               string\n\t}{\n\t\t{\n\t\t\tName:   \"no header\",\n\t\t\tServer: bytes.NewReader(nil),\n\t\t\tError:  \"EOF\",\n\t\t},\n\t\t{\n\t\t\tName: \"unknown member name\",\n\t\t\tServer: fakeSingleDocReply(\n\t\t\t\tmap[string]interface{}{\n\t\t\t\t\t\"members\": []map[string]interface{}{\n\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\"name\": \"foo\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t),\n\t\t\tError:               errProxyNotFound.Error(),\n\t\t\tProxyMapper:         fakeProxyMapper{},\n\t\t\tReplicaStateCompare: fakeReplicaStateCompare{sameIM: true, sameRS: true},\n\t\t},\n\t\t{\n\t\t\tName:                \"diffferent rs\",\n\t\t\tServer:              fakeSingleDocReply(map[string]interface{}{}),\n\t\t\tError:               errRSChanged.Error(),\n\t\t\tProxyMapper:         nil,\n\t\t\tReplicaStateCompare: fakeReplicaStateCompare{sameIM: true, sameRS: false},\n\t\t},\n\t}\n\n\tfor _, c := range cases {\n\t\tr := &ReplSetGetStatusResponseRewriter{\n\t\t\tLog:                 &tLogger{TB: t},\n\t\t\tProxyMapper:         c.ProxyMapper,\n\t\t\tReplicaStateCompare: c.ReplicaStateCompare,\n\t\t\tReplyRW: &ReplyRW{\n\t\t\t\tLog: &tLogger{TB: t},\n\t\t\t},\n\t\t}\n\t\terr := r.Rewrite(c.Client, c.Server)\n\t\tif err == nil {\n\t\t\tt.Errorf(\"was expecting an error for case %s\", c.Name)\n\t\t}\n\t\tif !strings.Contains(err.Error(), c.Error) {\n\t\t\tt.Errorf(\"did not get expected error for case %s instead got %s\", c.Name, err)\n\t\t}\n\t}\n}\n\nfunc TestReplSetGetStatusResponseRewriterSuccess(t *testing.T) {\n\tproxyMapper := fakeProxyMapper{\n\t\tm: map[string]string{\n\t\t\t\"a\": \"1\",\n\t\t\t\"b\": \"2\",\n\t\t\t\"c\": \"3\",\n\t\t},\n\t}\n\tin := bson.M{\n\t\t\"members\": []interface{}{\n\t\t\tbson.M{\n\t\t\t\t\"name\":     \"a\",\n\t\t\t\t\"stateStr\": \"PRIMARY\",\n\t\t\t},\n\t\t\tbson.M{\n\t\t\t\t\"name\": \"b\",\n\t\t\t},\n\t\t\tbson.M{\n\t\t\t\t\"name\":     \"c\",\n\t\t\t\t\"stateStr\": \"ARBITER\",\n\t\t\t},\n\t\t},\n\t}\n\tout := bson.M{\n\t\t\"members\": []interface{}{\n\t\t\tbson.M{\n\t\t\t\t\"name\":     \"1\",\n\t\t\t\t\"stateStr\": \"PRIMARY\",\n\t\t\t},\n\t\t\tbson.M{\n\t\t\t\t\"name\": \"2\",\n\t\t\t},\n\t\t\tbson.M{\n\t\t\t\t\"name\":     \"3\",\n\t\t\t\t\"stateStr\": \"ARBITER\",\n\t\t\t},\n\t\t},\n\t}\n\tr := &ReplSetGetStatusResponseRewriter{\n\t\tLog:                 &tLogger{TB: t},\n\t\tProxyMapper:         proxyMapper,\n\t\tReplicaStateCompare: fakeReplicaStateCompare{sameIM: true, sameRS: true},\n\t\tReplyRW: &ReplyRW{\n\t\t\tLog: &tLogger{TB: t},\n\t\t},\n\t}\n\n\tvar client bytes.Buffer\n\tif err := r.Rewrite(&client, fakeSingleDocReply(in)); err != nil {\n\t\tt.Fatal(err)\n\t}\n\tactualOut := bson.M{}\n\tdoc := client.Bytes()[headerLen+len(emptyPrefix):]\n\tif err := bson.Unmarshal(doc, &actualOut); err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif !reflect.DeepEqual(out, actualOut) {\n\t\tspew.Dump(out)\n\t\tspew.Dump(actualOut)\n\t\tt.Fatal(\"did not get expected output\")\n\t}\n}\n\nfunc TestProxyQuery(t *testing.T) {\n\tt.Parallel()\n\tvar p ProxyQuery\n\tlog := tLogger{TB: t}\n\tvar graph inject.Graph\n\terr := graph.Provide(\n\t\t&inject.Object{Value: &fakeProxyMapper{}},\n\t\t&inject.Object{Value: &fakeReplicaStateCompare{}},\n\t\t&inject.Object{Value: &log},\n\t\t&inject.Object{Value: &p},\n\t)\n\tensure.Nil(t, err)\n\tensure.Nil(t, graph.Populate())\n\tobjects := graph.Objects()\n\tensure.Nil(t, startstop.Start(objects, &log))\n\tdefer startstop.Stop(objects, &log)\n\n\tcases := []struct {\n\t\tName   string\n\t\tHeader *messageHeader\n\t\tClient io.ReadWriter\n\t\tError  string\n\t}{\n\t\t{\n\t\t\tName:   \"EOF while reading flags from client\",\n\t\t\tHeader: &messageHeader{},\n\t\t\tClient: new(bytes.Buffer),\n\t\t\tError:  \"EOF\",\n\t\t},\n\t\t{\n\t\t\tName:   \"EOF while reading collection name\",\n\t\t\tHeader: &messageHeader{},\n\t\t\tClient: fakeReadWriter{\n\t\t\t\tReader: bytes.NewReader(\n\t\t\t\t\t[]byte{0, 0, 0, 0}, // flags int32 before collection name\n\t\t\t\t),\n\t\t\t},\n\t\t\tError: \"EOF\",\n\t\t},\n\t\t{\n\t\t\tName:   \"EOF while reading skip/return\",\n\t\t\tHeader: &messageHeader{},\n\t\t\tClient: fakeReadWriter{\n\t\t\t\tReader: bytes.NewReader(\n\t\t\t\t\tappend(\n\t\t\t\t\t\t[]byte{0, 0, 0, 0}, // flags int32 before collection name\n\t\t\t\t\t\tadminCollectionName...,\n\t\t\t\t\t),\n\t\t\t\t),\n\t\t\t},\n\t\t\tError: \"EOF\",\n\t\t},\n\t\t{\n\t\t\tName:   \"EOF while reading query document\",\n\t\t\tHeader: &messageHeader{},\n\t\t\tClient: fakeReadWriter{\n\t\t\t\tReader: io.MultiReader(\n\t\t\t\t\tbytes.NewReader([]byte{0, 0, 0, 0}), // flags int32 before collection name\n\t\t\t\t\tbytes.NewReader(adminCollectionName),\n\t\t\t\t\tbytes.NewReader(\n\t\t\t\t\t\t[]byte{\n\t\t\t\t\t\t\t0, 0, 0, 0, // numberToSkip int32\n\t\t\t\t\t\t\t0, 0, 0, 0, // numberToReturn int32\n\t\t\t\t\t\t\t1, // partial bson document length header\n\t\t\t\t\t\t}),\n\t\t\t\t),\n\t\t\t},\n\t\t\tError: \"EOF\",\n\t\t},\n\t\t{\n\t\t\tName:   \"error while unmarshaling query document\",\n\t\t\tHeader: &messageHeader{},\n\t\t\tClient: fakeReadWriter{\n\t\t\t\tReader: io.MultiReader(\n\t\t\t\t\tbytes.NewReader([]byte{0, 0, 0, 0}), // flags int32 before collection name\n\t\t\t\t\tbytes.NewReader(adminCollectionName),\n\t\t\t\t\tbytes.NewReader(\n\t\t\t\t\t\t[]byte{\n\t\t\t\t\t\t\t0, 0, 0, 0, // numberToSkip int32\n\t\t\t\t\t\t\t0, 0, 0, 0, // numberToReturn int32\n\t\t\t\t\t\t\t5, 0, 0, 0, // bson document length header\n\t\t\t\t\t\t\t1, // bson document\n\t\t\t\t\t\t}),\n\t\t\t\t),\n\t\t\t},\n\t\t\tError: \"Document is corrupted\",\n\t\t},\n\t}\n\n\tfor _, c := range cases {\n\t\terr := p.Proxy(c.Header, c.Client, nil, nil)\n\t\tif err == nil || !strings.Contains(err.Error(), c.Error) {\n\t\t\tt.Fatalf(\"did not find expected error for %s, instead found %s\", c.Name, err)\n\t\t}\n\t}\n}\n"
  },
  {
    "path": "rs_state.go",
    "content": "package dvara\n\nimport (\n\t\"fmt\"\n\t\"sort\"\n\t\"time\"\n\n\t\"github.com/davecgh/go-spew/spew\"\n\n\t\"gopkg.in/mgo.v2\"\n\t\"gopkg.in/mgo.v2/bson\"\n)\n\nconst errNotReplSet = \"not running with --replSet\"\n\n// ReplicaSetState is a snapshot of the RS configuration at some point in time.\ntype ReplicaSetState struct {\n\tlastRS     *replSetGetStatusResponse\n\tlastIM     *isMasterResponse\n\tsingleAddr string // this is only set when we're not running against a RS\n}\n\n// NewReplicaSetState creates a new ReplicaSetState using the given address.\nfunc NewReplicaSetState(addr string) (*ReplicaSetState, error) {\n\tinfo := &mgo.DialInfo{\n\t\tAddrs:   []string{addr},\n\t\tDirect:  true,\n\t\tTimeout: 5 * time.Second,\n\t}\n\tsession, err := mgo.DialWithInfo(info)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tsession.SetMode(mgo.Monotonic, true)\n\tsession.SetSyncTimeout(5 * time.Second)\n\tsession.SetSocketTimeout(5 * time.Second)\n\tdefer session.Close()\n\n\tvar r ReplicaSetState\n\tif r.lastRS, err = replSetGetStatus(session); err != nil {\n\t\t// This error indicates we're in Single Node Mode. That's okay.\n\t\tif err.Error() != errNotReplSet {\n\t\t\treturn nil, err\n\t\t}\n\t\tr.singleAddr = addr\n\t}\n\n\tif r.lastIM, err = isMaster(session); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif r.lastRS != nil && len(r.lastRS.Members) == 1 {\n\t\tn := r.lastRS.Members[0]\n\t\tif n.State != \"PRIMARY\" || n.State != \"SECONDARY\" {\n\t\t\treturn nil, fmt.Errorf(\"single node RS in bad state: %s\", spew.Sdump(r))\n\t\t}\n\t}\n\n\t// nodes starting up are invalid\n\tif r.lastRS != nil {\n\t\tfor _, member := range r.lastRS.Members {\n\t\t\tif member.Self && member.State == \"STARTUP\" {\n\t\t\t\treturn nil, fmt.Errorf(\"node is busy starting up: %s\", member.Name)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn &r, nil\n}\n\n// AssertEqual checks if the given ReplicaSetState equals this one. It returns\n// a rich error message including the entire state for easier debugging.\nfunc (r *ReplicaSetState) AssertEqual(o *ReplicaSetState) error {\n\tif r.Equal(o) {\n\t\treturn nil\n\t}\n\treturn fmt.Errorf(\n\t\t\"conflicting ReplicaSetState:\\n%s\\nVS\\n%s\",\n\t\tspew.Sdump(r),\n\t\tspew.Sdump(o),\n\t)\n}\n\n// Equal returns true if the given ReplicaSetState is the same as this one.\nfunc (r *ReplicaSetState) Equal(o *ReplicaSetState) bool {\n\treturn r.SameIM(o.lastIM) && r.SameRS(o.lastRS)\n}\n\n// SameRS checks if the given replSetGetStatusResponse is the same as the one\n// we have.\nfunc (r *ReplicaSetState) SameRS(o *replSetGetStatusResponse) bool {\n\treturn sameRSMembers(r.lastRS, o)\n}\n\n// SameIM checks if the given isMasterResponse is the same as the one we have.\nfunc (r *ReplicaSetState) SameIM(o *isMasterResponse) bool {\n\treturn sameIMMembers(r.lastIM, o)\n}\n\n// Addrs returns the addresses of members in primary or secondary state.\nfunc (r *ReplicaSetState) Addrs() []string {\n\tif r.singleAddr != \"\" {\n\t\treturn []string{r.singleAddr}\n\t}\n\tvar members []string\n\tfor _, m := range r.lastRS.Members {\n\t\tif m.State == ReplicaStatePrimary || m.State == ReplicaStateSecondary {\n\t\t\tmembers = append(members, m.Name)\n\t\t}\n\t}\n\treturn members\n}\n\n// ReplicaSetStateCreator allows for creating a ReplicaSetState from a given\n// set of seed addresses.\ntype ReplicaSetStateCreator struct {\n\tLog Logger `inject:\"\"`\n}\n\n// FromAddrs creates a ReplicaSetState from the given set of see addresses. It\n// requires the addresses to be part of the same Replica Set.\nfunc (c *ReplicaSetStateCreator) FromAddrs(addrs []string, replicaSetName string) (*ReplicaSetState, error) {\n\tvar r *ReplicaSetState\n\tfor _, addr := range addrs {\n\t\tar, err := NewReplicaSetState(addr)\n\t\tif err != nil {\n\t\t\tc.Log.Errorf(\"ignoring failure against address %s: %s\", addr, err)\n\t\t\tcontinue\n\t\t}\n\n\t\tif replicaSetName != \"\" {\n\t\t\tif ar.lastRS == nil {\n\t\t\t\tc.Log.Errorf(\n\t\t\t\t\t\"ignoring standalone node %q not in expected replset: %q\",\n\t\t\t\t\taddr,\n\t\t\t\t\treplicaSetName,\n\t\t\t\t)\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tif ar.lastRS.Name != replicaSetName {\n\t\t\t\tc.Log.Errorf(\n\t\t\t\t\t\"ignoring node %q not in expected replset: %q vs %q\",\n\t\t\t\t\taddr,\n\t\t\t\t\tar.lastRS.Name,\n\t\t\t\t\treplicaSetName,\n\t\t\t\t)\n\t\t\t\tcontinue\n\t\t\t}\n\t\t}\n\n\t\t// First successful address.\n\t\tif r == nil {\n\t\t\tr = ar\n\t\t\tcontinue\n\t\t}\n\n\t\t// Ensure same as already established ReplicaSetState.\n\t\tif err := r.AssertEqual(ar); err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\tif r == nil {\n\t\treturn nil, fmt.Errorf(\"could not connect to any provided addresses: %v\", addrs)\n\t}\n\n\t// Check if we're expecting an RS but got a single node.\n\tif r.singleAddr != \"\" && len(addrs) != 1 {\n\t\treturn nil, fmt.Errorf(\n\t\t\t\"node %s is not in a replica set but was expecting it to be in a\"+\n\t\t\t\t\" replica set with members %v\",\n\t\t\tr.singleAddr,\n\t\t\taddrs,\n\t\t)\n\t}\n\n\treturn r, nil\n}\n\nvar (\n\treplSetGetStatusQuery = bson.D{\n\t\tbson.DocElem{Name: \"replSetGetStatus\", Value: 1},\n\t}\n\tisMasterQuery = bson.D{\n\t\tbson.DocElem{Name: \"isMaster\", Value: 1},\n\t}\n)\n\nfunc replSetGetStatus(s *mgo.Session) (*replSetGetStatusResponse, error) {\n\tvar res replSetGetStatusResponse\n\tif err := s.Run(replSetGetStatusQuery, &res); err != nil {\n\t\treturn nil, err\n\t}\n\treturn &res, nil\n}\n\nfunc isMaster(s *mgo.Session) (*isMasterResponse, error) {\n\tvar res isMasterResponse\n\tif err := s.Run(isMasterQuery, &res); err != nil {\n\t\treturn nil, fmt.Errorf(\"error in isMaster: %s\", err)\n\t}\n\treturn &res, nil\n}\n\nfunc sameRSMembers(a *replSetGetStatusResponse, b *replSetGetStatusResponse) bool {\n\tif (a == nil || len(a.Members) == 0) && (b == nil || len(b.Members) == 0) {\n\t\treturn true\n\t}\n\tif a == nil || b == nil {\n\t\treturn false\n\t}\n\tl := len(a.Members)\n\tif l != len(b.Members) {\n\t\treturn false\n\t}\n\taMembers := make([]string, 0, l)\n\tbMembers := make([]string, 0, l)\n\tfor i := 0; i < l; i++ {\n\t\taM := a.Members[i]\n\t\taMembers = append(aMembers, fmt.Sprintf(\"%s:%s\", aM.Name, aM.State))\n\t\tbM := b.Members[i]\n\t\tbMembers = append(bMembers, fmt.Sprintf(\"%s:%s\", bM.Name, bM.State))\n\t}\n\tsort.Strings(aMembers)\n\tsort.Strings(bMembers)\n\tfor i := 0; i < l; i++ {\n\t\tif aMembers[i] != bMembers[i] {\n\t\t\treturn false\n\t\t}\n\t}\n\treturn true\n}\n\nvar emptyIsMasterResponse = isMasterResponse{}\n\nfunc sameIMMembers(a *isMasterResponse, b *isMasterResponse) bool {\n\tif a == nil && b == nil {\n\t\treturn true\n\t}\n\tif a == nil {\n\t\ta = &emptyIsMasterResponse\n\t}\n\tif b == nil {\n\t\tb = &emptyIsMasterResponse\n\t}\n\tl := len(a.Hosts)\n\tif l != len(b.Hosts) {\n\t\treturn false\n\t}\n\taHosts := make([]string, 0, l+1)\n\tbHosts := make([]string, 0, l+1)\n\tfor i := 0; i < l; i++ {\n\t\taHosts = append(aHosts, a.Hosts[i])\n\t\tbHosts = append(bHosts, b.Hosts[i])\n\t}\n\tsort.Strings(aHosts)\n\tsort.Strings(bHosts)\n\taHosts = append(aHosts, a.Primary)\n\tbHosts = append(bHosts, b.Primary)\n\tfor i := range aHosts {\n\t\tif aHosts[i] != bHosts[i] {\n\t\t\treturn false\n\t\t}\n\t}\n\treturn true\n}\n"
  },
  {
    "path": "rs_state_test.go",
    "content": "package dvara\n\nimport (\n\t\"testing\"\n\n\t\"github.com/facebookgo/mgotest\"\n)\n\nfunc TestSameRSMembers(t *testing.T) {\n\tt.Parallel()\n\tcases := []struct {\n\t\tName string\n\t\tA    *replSetGetStatusResponse\n\t\tB    *replSetGetStatusResponse\n\t}{\n\t\t{\n\t\t\tName: \"the same\",\n\t\t\tA: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"a\", State: \"b\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\tB: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"a\", State: \"b\"},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"out of order\",\n\t\t\tA: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"a\", State: \"b\"},\n\t\t\t\t\t{Name: \"c\", State: \"d\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\tB: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"c\", State: \"d\"},\n\t\t\t\t\t{Name: \"a\", State: \"b\"},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"both nil\",\n\t\t},\n\t\t{\n\t\t\tName: \"A nil B empty\",\n\t\t\tB:    &replSetGetStatusResponse{},\n\t\t},\n\t\t{\n\t\t\tName: \"A empty B nil\",\n\t\t\tA:    &replSetGetStatusResponse{},\n\t\t},\n\t}\n\n\tfor _, c := range cases {\n\t\tif !sameRSMembers(c.A, c.B) {\n\t\t\tt.Fatalf(\"failed %s\", c.Name)\n\t\t}\n\t}\n}\n\nfunc TestNotSameRSMembers(t *testing.T) {\n\tt.Parallel()\n\tcases := []struct {\n\t\tName string\n\t\tA    *replSetGetStatusResponse\n\t\tB    *replSetGetStatusResponse\n\t}{\n\t\t{\n\t\t\tName: \"different name\",\n\t\t\tA: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"a\", State: \"b\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\tB: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"b\", State: \"b\"},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"different state\",\n\t\t\tA: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"a\", State: \"b\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\tB: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"a\", State: \"c\"},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"subset A\",\n\t\t\tA: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"a\", State: \"b\"},\n\t\t\t\t\t{Name: \"b\", State: \"c\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\tB: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"a\", State: \"b\"},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"subset B\",\n\t\t\tA: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"a\", State: \"b\"},\n\t\t\t\t},\n\t\t\t},\n\t\t\tB: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"a\", State: \"b\"},\n\t\t\t\t\t{Name: \"b\", State: \"c\"},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"nil A\",\n\t\t\tB: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"b\", State: \"b\"},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"nil B\",\n\t\t\tA: &replSetGetStatusResponse{\n\t\t\t\tMembers: []statusMember{\n\t\t\t\t\t{Name: \"a\", State: \"b\"},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}\n\n\tfor _, c := range cases {\n\t\tif sameRSMembers(c.A, c.B) {\n\t\t\tt.Fatalf(\"failed %s\", c.Name)\n\t\t}\n\t}\n}\n\nfunc TestSameIMMembers(t *testing.T) {\n\tt.Parallel()\n\tcases := []struct {\n\t\tName string\n\t\tA    *isMasterResponse\n\t\tB    *isMasterResponse\n\t}{\n\t\t{\n\t\t\tName: \"the same\",\n\t\t\tA: &isMasterResponse{\n\t\t\t\tHosts: []string{\"a\", \"b\"},\n\t\t\t},\n\t\t\tB: &isMasterResponse{\n\t\t\t\tHosts: []string{\"a\", \"b\"},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"out of order\",\n\t\t\tA: &isMasterResponse{\n\t\t\t\tHosts: []string{\"a\", \"b\"},\n\t\t\t},\n\t\t\tB: &isMasterResponse{\n\t\t\t\tHosts: []string{\"b\", \"a\"},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"both nil\",\n\t\t},\n\t\t{\n\t\t\tName: \"A nil B empty\",\n\t\t\tB:    &isMasterResponse{},\n\t\t},\n\t\t{\n\t\t\tName: \"A empty B nil\",\n\t\t\tA:    &isMasterResponse{},\n\t\t},\n\t}\n\n\tfor _, c := range cases {\n\t\tif !sameIMMembers(c.A, c.B) {\n\t\t\tt.Fatalf(\"failed %s\", c.Name)\n\t\t}\n\t}\n}\n\nfunc TestNotSameIMMembers(t *testing.T) {\n\tt.Parallel()\n\tcases := []struct {\n\t\tName string\n\t\tA    *isMasterResponse\n\t\tB    *isMasterResponse\n\t}{\n\t\t{\n\t\t\tName: \"different name\",\n\t\t\tA: &isMasterResponse{\n\t\t\t\tHosts: []string{\"a\"},\n\t\t\t},\n\t\t\tB: &isMasterResponse{\n\t\t\t\tHosts: []string{\"b\"},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"subset A\",\n\t\t\tA: &isMasterResponse{\n\t\t\t\tHosts: []string{\"a\", \"b\"},\n\t\t\t},\n\t\t\tB: &isMasterResponse{\n\t\t\t\tHosts: []string{\"a\"},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"subset B\",\n\t\t\tA: &isMasterResponse{\n\t\t\t\tHosts: []string{\"a\"},\n\t\t\t},\n\t\t\tB: &isMasterResponse{\n\t\t\t\tHosts: []string{\"a\", \"b\"},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"nil A\",\n\t\t\tB: &isMasterResponse{\n\t\t\t\tHosts: []string{\"a\"},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tName: \"nil B\",\n\t\t\tA: &isMasterResponse{\n\t\t\t\tHosts: []string{\"b\"},\n\t\t\t},\n\t\t},\n\t}\n\n\tfor _, c := range cases {\n\t\tif sameIMMembers(c.A, c.B) {\n\t\t\tt.Fatalf(\"failed %s\", c.Name)\n\t\t}\n\t}\n}\n\nfunc TestSingleNodeNewReplicaSetState(t *testing.T) {\n\tt.Parallel()\n\tmgo := mgotest.NewStartedServer(t)\n\tdefer mgo.Stop()\n\trs, err := NewReplicaSetState(mgo.URL())\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\tif rs.singleAddr != mgo.URL() {\n\t\tt.Fatalf(\"expected %s got %s\", mgo.URL(), rs.singleAddr)\n\t}\n}\n\nfunc TestNewReplicaSetStateFailure(t *testing.T) {\n\tt.Parallel()\n\tmgo := mgotest.NewStartedServer(t)\n\tmgo.Stop()\n\t_, err := NewReplicaSetState(mgo.URL())\n\tconst expected = \"no reachable servers\"\n\tif err == nil || err.Error() != expected {\n\t\tt.Fatalf(\"unexpected error: %s\", err)\n\t}\n}\n\nfunc TestSingleNodeNewReplicaSetStateAddrs(t *testing.T) {\n\tt.Parallel()\n\tmgo := mgotest.NewStartedServer(t)\n\tdefer mgo.Stop()\n\trs, err := NewReplicaSetState(mgo.URL())\n\tif err != nil {\n\t\tt.Fatal(err)\n\t}\n\taddrs := rs.Addrs()\n\tif len(addrs) != 1 || addrs[0] != mgo.URL() {\n\t\tt.Fatalf(\"unexpected addrs %v\", addrs)\n\t}\n}\n\nfunc TestIgnoreMismatchingReplicaSets(t *testing.T) {\n\tif disableSlowTests {\n\t\tt.Skip(\"disabled because it's slow\")\n\t}\n\tt.Parallel()\n\tcreator := ReplicaSetStateCreator{\n\t\tLog: &tLogger{TB: t},\n\t}\n\treplicaSet := mgotest.NewReplicaSet(2, t)\n\tsingleMongo := mgotest.NewStartedServer(t)\n\tdefer func() {\n\t\treplicaSet.Stop()\n\t\tsingleMongo.Stop()\n\t}()\n\n\turls := replicaSet.Addrs()\n\turls = append(urls, singleMongo.URL())\n\n\tstate, err := creator.FromAddrs(urls, \"rs\")\n\tif err != nil {\n\t\tt.Fatalf(\"unexpected error: %s\", err)\n\t}\n\tif state.lastRS.Name != \"rs\" {\n\t\tt.Fatalf(\"unexpected replicaset: %s\", state.lastRS.Name)\n\t}\n\n\t_, err = creator.FromAddrs(urls, \"\")\n\tif err == nil {\n\t\tt.Fatalf(\"missing expected error: %s\", err)\n\t}\n}\n"
  },
  {
    "path": "state.go",
    "content": "package dvara\n\n// ReplicaState is the state of a node in the replica.\ntype ReplicaState string\n\nconst (\n\t// ReplicaStatePrimary indicates the node is a primary.\n\tReplicaStatePrimary = ReplicaState(\"PRIMARY\")\n\n\t// ReplicaStateSecondary indicates the node is a secondary.\n\tReplicaStateSecondary = ReplicaState(\"SECONDARY\")\n\n\t// ReplicaStateArbiter indicates the node is an arbiter.\n\tReplicaStateArbiter = ReplicaState(\"ARBITER\")\n)\n"
  }
]