← Back to branch summary

~nskaggs/+junk/xenial-test

« back to all changes in this revision

Viewing changes to src/gopkg.in/mgo.v2/dbtest/dbserver.go

  • Committer: Nicholas Skaggs
  • Date: 2016-10-24 20:56:05 UTC
  • Revision ID: nicholas.skaggs@canonical.com-20161024205605-z8lta0uvuhtxwzwl
Initi with beta15

Show diffs side-by-side

added added

removed removed

Lines of Context:
src/gopkg.in/mgo.v2/dbtest/dbserver.go
 
1
package dbtest
 
2
 
 
3
import (
 
4
        "bytes"
 
5
        "fmt"
 
6
        "net"
 
7
        "os"
 
8
        "os/exec"
 
9
        "strconv"
 
10
        "time"
 
11
 
 
12
        "gopkg.in/mgo.v2"
 
13
        "gopkg.in/tomb.v2"
 
14
)
 
15
 
 
16
// DBServer controls a MongoDB server process to be used within test suites.
 
17
//
 
18
// The test server is started when Session is called the first time and should
 
19
// remain running for the duration of all tests, with the Wipe method being
 
20
// called between tests (before each of them) to clear stored data. After all tests
 
21
// are done, the Stop method should be called to stop the test server.
 
22
//
 
23
// Before the DBServer is used the SetPath method must be called to define
 
24
// the location for the database files to be stored.
 
25
type DBServer struct {
 
26
        session *mgo.Session
 
27
        output  bytes.Buffer
 
28
        server  *exec.Cmd
 
29
        dbpath  string
 
30
        host    string
 
31
        tomb    tomb.Tomb
 
32
}
 
33
 
 
34
// SetPath defines the path to the directory where the database files will be
 
35
// stored if it is started. The directory path itself is not created or removed
 
36
// by the test helper.
 
37
func (dbs *DBServer) SetPath(dbpath string) {
 
38
        dbs.dbpath = dbpath
 
39
}
 
40
 
 
41
func (dbs *DBServer) start() {
 
42
        if dbs.server != nil {
 
43
                panic("DBServer already started")
 
44
        }
 
45
        if dbs.dbpath == "" {
 
46
                panic("DBServer.SetPath must be called before using the server")
 
47
        }
 
48
        mgo.SetStats(true)
 
49
        l, err := net.Listen("tcp", "127.0.0.1:0")
 
50
        if err != nil {
 
51
                panic("unable to listen on a local address: " + err.Error())
 
52
        }
 
53
        addr := l.Addr().(*net.TCPAddr)
 
54
        l.Close()
 
55
        dbs.host = addr.String()
 
56
 
 
57
        args := []string{
 
58
                "--dbpath", dbs.dbpath,
 
59
                "--bind_ip", "127.0.0.1",
 
60
                "--port", strconv.Itoa(addr.Port),
 
61
                "--nssize", "1",
 
62
                "--noprealloc",
 
63
                "--smallfiles",
 
64
                "--nojournal",
 
65
        }
 
66
        dbs.tomb = tomb.Tomb{}
 
67
        dbs.server = exec.Command("mongod", args...)
 
68
        dbs.server.Stdout = &dbs.output
 
69
        dbs.server.Stderr = &dbs.output
 
70
        err = dbs.server.Start()
 
71
        if err != nil {
 
72
                panic(err)
 
73
        }
 
74
        dbs.tomb.Go(dbs.monitor)
 
75
        dbs.Wipe()
 
76
}
 
77
 
 
78
func (dbs *DBServer) monitor() error {
 
79
        dbs.server.Process.Wait()
 
80
        if dbs.tomb.Alive() {
 
81
                // Present some debugging information.
 
82
                fmt.Fprintf(os.Stderr, "---- mongod process died unexpectedly:\n")
 
83
                fmt.Fprintf(os.Stderr, "%s", dbs.output.Bytes())
 
84
                fmt.Fprintf(os.Stderr, "---- mongod processes running right now:\n")
 
85
                cmd := exec.Command("/bin/sh", "-c", "ps auxw | grep mongod")
 
86
                cmd.Stdout = os.Stderr
 
87
                cmd.Stderr = os.Stderr
 
88
                cmd.Run()
 
89
                fmt.Fprintf(os.Stderr, "----------------------------------------\n")
 
90
 
 
91
                panic("mongod process died unexpectedly")
 
92
        }
 
93
        return nil
 
94
}
 
95
 
 
96
// Stop stops the test server process, if it is running.
 
97
//
 
98
// It's okay to call Stop multiple times. After the test server is
 
99
// stopped it cannot be restarted.
 
100
//
 
101
// All database sessions must be closed before or while the Stop method
 
102
// is running. Otherwise Stop will panic after a timeout informing that
 
103
// there is a session leak.
 
104
func (dbs *DBServer) Stop() {
 
105
        if dbs.session != nil {
 
106
                dbs.checkSessions()
 
107
                if dbs.session != nil {
 
108
                        dbs.session.Close()
 
109
                        dbs.session = nil
 
110
                }
 
111
        }
 
112
        if dbs.server != nil {
 
113
                dbs.tomb.Kill(nil)
 
114
                dbs.server.Process.Signal(os.Interrupt)
 
115
                select {
 
116
                case <-dbs.tomb.Dead():
 
117
                case <-time.After(5 * time.Second):
 
118
                        panic("timeout waiting for mongod process to die")
 
119
                }
 
120
                dbs.server = nil
 
121
        }
 
122
}
 
123
 
 
124
// Session returns a new session to the server. The returned session
 
125
// must be closed after the test is done with it.
 
126
//
 
127
// The first Session obtained from a DBServer will start it.
 
128
func (dbs *DBServer) Session() *mgo.Session {
 
129
        if dbs.server == nil {
 
130
                dbs.start()
 
131
        }
 
132
        if dbs.session == nil {
 
133
                mgo.ResetStats()
 
134
                var err error
 
135
                dbs.session, err = mgo.Dial(dbs.host + "/test")
 
136
                if err != nil {
 
137
                        panic(err)
 
138
                }
 
139
        }
 
140
        return dbs.session.Copy()
 
141
}
 
142
 
 
143
// checkSessions ensures all mgo sessions opened were properly closed.
 
144
// For slightly faster tests, it may be disabled setting the
 
145
// environmnet variable CHECK_SESSIONS to 0.
 
146
func (dbs *DBServer) checkSessions() {
 
147
        if check := os.Getenv("CHECK_SESSIONS"); check == "0" || dbs.server == nil || dbs.session == nil {
 
148
                return
 
149
        }
 
150
        dbs.session.Close()
 
151
        dbs.session = nil
 
152
        for i := 0; i < 100; i++ {
 
153
                stats := mgo.GetStats()
 
154
                if stats.SocketsInUse == 0 && stats.SocketsAlive == 0 {
 
155
                        return
 
156
                }
 
157
                time.Sleep(100 * time.Millisecond)
 
158
        }
 
159
        panic("There are mgo sessions still alive.")
 
160
}
 
161
 
 
162
// Wipe drops all created databases and their data.
 
163
//
 
164
// The MongoDB server remains running if it was prevoiusly running,
 
165
// or stopped if it was previously stopped.
 
166
//
 
167
// All database sessions must be closed before or while the Wipe method
 
168
// is running. Otherwise Wipe will panic after a timeout informing that
 
169
// there is a session leak.
 
170
func (dbs *DBServer) Wipe() {
 
171
        if dbs.server == nil || dbs.session == nil {
 
172
                return
 
173
        }
 
174
        dbs.checkSessions()
 
175
        sessionUnset := dbs.session == nil
 
176
        session := dbs.Session()
 
177
        defer session.Close()
 
178
        if sessionUnset {
 
179
                dbs.session.Close()
 
180
                dbs.session = nil
 
181
        }
 
182
        names, err := session.DatabaseNames()
 
183
        if err != nil {
 
184
                panic(err)
 
185
        }
 
186
        for _, name := range names {
 
187
                switch name {
 
188
                case "admin", "local", "config":
 
189
                default:
 
190
                        err = session.DB(name).DropDatabase()
 
191
                        if err != nil {
 
192
                                panic(err)
 
193
                        }
 
194
                }
 
195
        }
 
196
}