462 lines
14 KiB
Python
462 lines
14 KiB
Python
# Copyright 2009-2012 10gen, Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Test that pymongo is thread safe."""
|
|
|
|
import unittest
|
|
import threading
|
|
import traceback
|
|
|
|
from nose.plugins.skip import SkipTest
|
|
|
|
from test.utils import server_started_with_auth, joinall
|
|
from test.test_connection import get_connection
|
|
from pymongo.connection import Connection
|
|
from pymongo.replica_set_connection import ReplicaSetConnection
|
|
from pymongo.pool import SocketInfo, _closed
|
|
from pymongo.errors import (AutoReconnect,
|
|
OperationFailure,
|
|
DuplicateKeyError)
|
|
|
|
|
|
def get_pool(connection):
|
|
if isinstance(connection, Connection):
|
|
return connection._Connection__pool
|
|
elif isinstance(connection, ReplicaSetConnection):
|
|
writer = connection._ReplicaSetConnection__writer
|
|
pools = connection._ReplicaSetConnection__pools
|
|
return pools[writer]['pool']
|
|
else:
|
|
raise TypeError(str(connection))
|
|
|
|
|
|
class AutoAuthenticateThreads(threading.Thread):
|
|
|
|
def __init__(self, collection, num):
|
|
threading.Thread.__init__(self)
|
|
self.coll = collection
|
|
self.num = num
|
|
self.success = True
|
|
self.setDaemon(True)
|
|
|
|
def run(self):
|
|
try:
|
|
for i in xrange(self.num):
|
|
self.coll.insert({'num':i}, safe=True)
|
|
self.coll.find_one({'num':i})
|
|
except Exception:
|
|
traceback.print_exc()
|
|
self.success = False
|
|
|
|
|
|
class SaveAndFind(threading.Thread):
|
|
|
|
def __init__(self, collection):
|
|
threading.Thread.__init__(self)
|
|
self.collection = collection
|
|
self.setDaemon(True)
|
|
|
|
def run(self):
|
|
sum = 0
|
|
for document in self.collection.find():
|
|
sum += document["x"]
|
|
assert sum == 499500, "sum was %d not 499500" % sum
|
|
|
|
|
|
class Insert(threading.Thread):
|
|
|
|
def __init__(self, collection, n, expect_exception):
|
|
threading.Thread.__init__(self)
|
|
self.collection = collection
|
|
self.n = n
|
|
self.expect_exception = expect_exception
|
|
self.setDaemon(True)
|
|
|
|
def run(self):
|
|
for _ in xrange(self.n):
|
|
error = True
|
|
|
|
try:
|
|
self.collection.insert({"test": "insert"}, safe=True)
|
|
error = False
|
|
except:
|
|
if not self.expect_exception:
|
|
raise
|
|
|
|
if self.expect_exception:
|
|
assert error
|
|
|
|
|
|
class Update(threading.Thread):
|
|
|
|
def __init__(self, collection, n, expect_exception):
|
|
threading.Thread.__init__(self)
|
|
self.collection = collection
|
|
self.n = n
|
|
self.expect_exception = expect_exception
|
|
self.setDaemon(True)
|
|
|
|
def run(self):
|
|
for _ in xrange(self.n):
|
|
error = True
|
|
|
|
try:
|
|
self.collection.update({"test": "unique"},
|
|
{"$set": {"test": "update"}}, safe=True)
|
|
error = False
|
|
except:
|
|
if not self.expect_exception:
|
|
raise
|
|
|
|
if self.expect_exception:
|
|
assert error
|
|
|
|
|
|
class IgnoreAutoReconnect(threading.Thread):
|
|
|
|
def __init__(self, collection, n):
|
|
threading.Thread.__init__(self)
|
|
self.c = collection
|
|
self.n = n
|
|
self.setDaemon(True)
|
|
|
|
def run(self):
|
|
for _ in range(self.n):
|
|
try:
|
|
self.c.find_one()
|
|
except AutoReconnect:
|
|
pass
|
|
|
|
|
|
class FindPauseFind(threading.Thread):
|
|
"""See test_server_disconnect() for details"""
|
|
@classmethod
|
|
def shared_state(cls, nthreads):
|
|
class SharedState(object):
|
|
pass
|
|
|
|
state = SharedState()
|
|
|
|
# Number of threads total
|
|
state.nthreads = nthreads
|
|
|
|
# Number of threads that have arrived at rendezvous point
|
|
state.arrived_threads = 0
|
|
state.arrived_threads_lock = threading.Lock()
|
|
|
|
# set when all threads reach rendezvous
|
|
state.ev_arrived = threading.Event()
|
|
|
|
# set from outside FindPauseFind to let threads resume after
|
|
# rendezvous
|
|
state.ev_resume = threading.Event()
|
|
return state
|
|
|
|
def __init__(self, collection, state):
|
|
"""Params: A collection, an event to signal when all threads have
|
|
done the first find(), an event to signal when threads should resume,
|
|
and the total number of threads
|
|
"""
|
|
super(FindPauseFind, self).__init__()
|
|
self.collection = collection
|
|
self.state = state
|
|
self.passed = False
|
|
|
|
# If this thread fails to terminate, don't hang the whole program
|
|
self.setDaemon(True)
|
|
|
|
def rendezvous(self):
|
|
# pause until all threads arrive here
|
|
s = self.state
|
|
s.arrived_threads_lock.acquire()
|
|
s.arrived_threads += 1
|
|
if s.arrived_threads == s.nthreads:
|
|
s.arrived_threads_lock.release()
|
|
s.ev_arrived.set()
|
|
else:
|
|
s.arrived_threads_lock.release()
|
|
s.ev_arrived.wait()
|
|
|
|
def run(self):
|
|
try:
|
|
# acquire a socket
|
|
list(self.collection.find())
|
|
|
|
pool = get_pool(self.collection.database.connection)
|
|
socket_info = pool._get_request_state()
|
|
assert isinstance(socket_info, SocketInfo)
|
|
self.request_sock = socket_info.sock
|
|
assert not _closed(self.request_sock)
|
|
|
|
# Dereference socket_info so it can potentially return to the pool
|
|
del socket_info
|
|
finally:
|
|
self.rendezvous()
|
|
|
|
# all threads have passed the rendezvous, wait for
|
|
# test_server_disconnect() to disconnect the connection
|
|
self.state.ev_resume.wait()
|
|
|
|
# test_server_disconnect() has closed this socket, but that's ok
|
|
# because it's not our request socket anymore
|
|
assert _closed(self.request_sock)
|
|
|
|
# if disconnect() properly closed all threads' request sockets, then
|
|
# this won't raise AutoReconnect because it will acquire a new socket
|
|
assert self.request_sock == pool._get_request_state().sock
|
|
list(self.collection.find())
|
|
assert self.collection.database.connection.in_request()
|
|
assert self.request_sock != pool._get_request_state().sock
|
|
self.passed = True
|
|
|
|
|
|
class BaseTestThreads(object):
|
|
"""
|
|
Base test class for TestThreads and TestThreadsReplicaSet. (This is not
|
|
itself a unittest.TestCase, otherwise it'd be run twice -- once when nose
|
|
imports this module, and once when nose imports
|
|
test_threads_replica_set_connection.py, which imports this module.)
|
|
"""
|
|
def setUp(self):
|
|
self.db = self._get_connection().pymongo_test
|
|
|
|
def tearDown(self):
|
|
pass
|
|
|
|
def _get_connection(self):
|
|
"""
|
|
Intended for overriding in TestThreadsReplicaSet. This method
|
|
returns a Connection here, and a ReplicaSetConnection in
|
|
test_threads_replica_set_connection.py.
|
|
"""
|
|
# Regular test connection
|
|
return get_connection()
|
|
|
|
def test_threading(self):
|
|
self.db.drop_collection("test")
|
|
for i in xrange(1000):
|
|
self.db.test.save({"x": i}, safe=True)
|
|
|
|
threads = []
|
|
for i in range(10):
|
|
t = SaveAndFind(self.db.test)
|
|
t.start()
|
|
threads.append(t)
|
|
|
|
joinall(threads)
|
|
|
|
def test_safe_insert(self):
|
|
self.db.drop_collection("test1")
|
|
self.db.test1.insert({"test": "insert"})
|
|
self.db.drop_collection("test2")
|
|
self.db.test2.insert({"test": "insert"})
|
|
|
|
self.db.test2.create_index("test", unique=True)
|
|
self.db.test2.find_one()
|
|
|
|
okay = Insert(self.db.test1, 2000, False)
|
|
error = Insert(self.db.test2, 2000, True)
|
|
|
|
error.start()
|
|
okay.start()
|
|
|
|
error.join()
|
|
okay.join()
|
|
|
|
def test_safe_update(self):
|
|
self.db.drop_collection("test1")
|
|
self.db.test1.insert({"test": "update"})
|
|
self.db.test1.insert({"test": "unique"})
|
|
self.db.drop_collection("test2")
|
|
self.db.test2.insert({"test": "update"})
|
|
self.db.test2.insert({"test": "unique"})
|
|
|
|
self.db.test2.create_index("test", unique=True)
|
|
self.db.test2.find_one()
|
|
|
|
okay = Update(self.db.test1, 2000, False)
|
|
error = Update(self.db.test2, 2000, True)
|
|
|
|
error.start()
|
|
okay.start()
|
|
|
|
error.join()
|
|
okay.join()
|
|
|
|
def test_low_network_timeout(self):
|
|
db = None
|
|
i = 0
|
|
n = 10
|
|
while db is None and i < n:
|
|
try:
|
|
db = get_connection(network_timeout=0.0001).pymongo_test
|
|
except AutoReconnect:
|
|
i += 1
|
|
if i == n:
|
|
raise SkipTest()
|
|
|
|
threads = []
|
|
for _ in range(4):
|
|
t = IgnoreAutoReconnect(db.test, 100)
|
|
t.start()
|
|
threads.append(t)
|
|
|
|
joinall(threads)
|
|
|
|
def test_server_disconnect(self):
|
|
# PYTHON-345, we need to make sure that threads' request sockets are
|
|
# closed by disconnect().
|
|
#
|
|
# 1. Create a connection with auto_start_request=True
|
|
# 2. Start N threads and do a find() in each to get a request socket
|
|
# 3. Pause all threads
|
|
# 4. In the main thread close all sockets, including threads' request
|
|
# sockets
|
|
# 5. In main thread, do a find(), which raises AutoReconnect and resets
|
|
# pool
|
|
# 6. Resume all threads, do a find() in them
|
|
#
|
|
# If we've fixed PYTHON-345, then only one AutoReconnect is raised,
|
|
# and all the threads get new request sockets.
|
|
|
|
cx = self.db.connection
|
|
self.assertTrue(cx.auto_start_request)
|
|
collection = self.db.pymongo_test
|
|
|
|
# acquire a request socket for the main thread
|
|
collection.find_one()
|
|
pool = get_pool(collection.database.connection)
|
|
socket_info = pool._get_request_state()
|
|
assert isinstance(socket_info, SocketInfo)
|
|
request_sock = socket_info.sock
|
|
|
|
state = FindPauseFind.shared_state(nthreads=40)
|
|
|
|
threads = [
|
|
FindPauseFind(collection, state)
|
|
for _ in range(state.nthreads)
|
|
]
|
|
|
|
# Each thread does a find(), thus acquiring a request socket
|
|
for t in threads:
|
|
t.start()
|
|
|
|
# Wait for the threads to reach the rendezvous
|
|
state.ev_arrived.wait(10)
|
|
self.assertTrue(state.ev_arrived.isSet(), "Thread timeout")
|
|
|
|
try:
|
|
self.assertEqual(state.nthreads, state.arrived_threads)
|
|
|
|
# Simulate an event that closes all sockets, e.g. primary stepdown
|
|
for t in threads:
|
|
t.request_sock.close()
|
|
|
|
# Finally, ensure the main thread's socket's last_checkout is
|
|
# updated:
|
|
collection.find_one()
|
|
|
|
# ... and close it:
|
|
request_sock.close()
|
|
|
|
# Doing an operation on the connection raises an AutoReconnect and
|
|
# resets the pool behind the scenes
|
|
self.assertRaises(AutoReconnect, collection.find_one)
|
|
|
|
finally:
|
|
# Let threads do a second find()
|
|
state.ev_resume.set()
|
|
|
|
joinall(threads)
|
|
|
|
for t in threads:
|
|
self.assertTrue(t.passed, "%s threw exception" % t)
|
|
|
|
|
|
class BaseTestThreadsAuth(object):
|
|
"""
|
|
Base test class for TestThreadsAuth and TestThreadsAuthReplicaSet. (This is
|
|
not itself a unittest.TestCase, otherwise it'd be run twice -- once when
|
|
nose imports this module, and once when nose imports
|
|
test_threads_replica_set_connection.py, which imports this module.)
|
|
"""
|
|
def _get_connection(self):
|
|
"""
|
|
Intended for overriding in TestThreadsAuthReplicaSet. This method
|
|
returns a Connection here, and a ReplicaSetConnection in
|
|
test_threads_replica_set_connection.py.
|
|
"""
|
|
# Regular test connection
|
|
return get_connection()
|
|
|
|
def setUp(self):
|
|
self.conn = self._get_connection()
|
|
if not server_started_with_auth(self.conn):
|
|
raise SkipTest("Authentication is not enabled on server")
|
|
self.conn.admin.system.users.remove({})
|
|
self.conn.admin.add_user('admin-user', 'password')
|
|
self.conn.admin.authenticate("admin-user", "password")
|
|
self.conn.auth_test.system.users.remove({})
|
|
self.conn.auth_test.add_user("test-user", "password")
|
|
|
|
def tearDown(self):
|
|
# Remove auth users from databases
|
|
self.conn.admin.authenticate("admin-user", "password")
|
|
self.conn.admin.system.users.remove({})
|
|
self.conn.auth_test.system.users.remove({})
|
|
|
|
def test_auto_auth_login(self):
|
|
conn = self._get_connection()
|
|
self.assertRaises(OperationFailure, conn.auth_test.test.find_one)
|
|
|
|
# Admin auth
|
|
conn = self._get_connection()
|
|
conn.admin.authenticate("admin-user", "password")
|
|
|
|
nthreads = 10
|
|
threads = []
|
|
for _ in xrange(nthreads):
|
|
t = AutoAuthenticateThreads(conn.auth_test.test, 100)
|
|
t.start()
|
|
threads.append(t)
|
|
|
|
joinall(threads)
|
|
|
|
for t in threads:
|
|
self.assertTrue(t.success)
|
|
|
|
# Database-specific auth
|
|
conn = self._get_connection()
|
|
conn.auth_test.authenticate("test-user", "password")
|
|
|
|
threads = []
|
|
for _ in xrange(nthreads):
|
|
t = AutoAuthenticateThreads(conn.auth_test.test, 100)
|
|
t.start()
|
|
threads.append(t)
|
|
|
|
joinall(threads)
|
|
|
|
for t in threads:
|
|
self.assertTrue(t.success)
|
|
|
|
class TestThreads(BaseTestThreads, unittest.TestCase):
|
|
pass
|
|
|
|
class TestThreadsAuth(BaseTestThreadsAuth, unittest.TestCase):
|
|
pass
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|