10a8c90248264a8b26970b4473770bcc3df8515fJosh Gao"""TestCases for distributed transactions.
20a8c90248264a8b26970b4473770bcc3df8515fJosh Gao"""
30a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
40a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoimport os
50a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoimport unittest
60a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
70a8c90248264a8b26970b4473770bcc3df8515fJosh Gaofrom test_all import db, test_support, get_new_environment_path, \
80a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        get_new_database_path
90a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
100a8c90248264a8b26970b4473770bcc3df8515fJosh Gaofrom test_all import verbose
110a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
120a8c90248264a8b26970b4473770bcc3df8515fJosh Gao#----------------------------------------------------------------------
130a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
140a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass DBTxn_distributed(unittest.TestCase):
150a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    num_txns=1234
160a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    nosync=True
170a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    must_open_db=False
180a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def _create_env(self, must_open_db) :
190a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.dbenv = db.DBEnv()
200a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.dbenv.set_tx_max(self.num_txns)
210a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.dbenv.set_lk_max_lockers(self.num_txns*2)
220a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.dbenv.set_lk_max_locks(self.num_txns*2)
230a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.dbenv.set_lk_max_objects(self.num_txns*2)
240a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if self.nosync :
250a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.dbenv.set_flags(db.DB_TXN_NOSYNC,True)
260a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.dbenv.open(self.homeDir, db.DB_CREATE | db.DB_THREAD |
270a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                db.DB_RECOVER |
280a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                db.DB_INIT_TXN | db.DB_INIT_LOG | db.DB_INIT_MPOOL |
290a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                db.DB_INIT_LOCK, 0666)
300a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.db = db.DB(self.dbenv)
310a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.db.set_re_len(db.DB_GID_SIZE)
320a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if must_open_db :
330a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            txn=self.dbenv.txn_begin()
340a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.db.open(self.filename,
350a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    db.DB_QUEUE, db.DB_CREATE | db.DB_THREAD, 0666,
360a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    txn=txn)
370a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            txn.commit()
380a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
390a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def setUp(self) :
400a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.homeDir = get_new_environment_path()
410a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.filename = "test"
420a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        return self._create_env(must_open_db=True)
430a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
440a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def _destroy_env(self):
450a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if self.nosync or (db.version()[:2] == (4,6)):  # Known bug
460a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.dbenv.log_flush()
470a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.db.close()
480a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.dbenv.close()
490a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
500a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def tearDown(self):
510a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self._destroy_env()
520a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        test_support.rmtree(self.homeDir)
530a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
540a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def _recreate_env(self,must_open_db) :
550a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self._destroy_env()
560a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self._create_env(must_open_db)
570a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
580a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def test01_distributed_transactions(self) :
590a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        txns=set()
600a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        adapt = lambda x : x
610a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        import sys
620a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if sys.version_info[0] >= 3 :
630a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            adapt = lambda x : bytes(x, "ascii")
640a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # Create transactions, "prepare" them, and
650a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # let them be garbage collected.
660a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for i in xrange(self.num_txns) :
670a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            txn = self.dbenv.txn_begin()
680a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            gid = "%%%dd" %db.DB_GID_SIZE
690a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            gid = adapt(gid %i)
700a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.db.put(i, gid, txn=txn, flags=db.DB_APPEND)
710a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            txns.add(gid)
720a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            txn.prepare(gid)
730a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        del txn
740a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
750a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self._recreate_env(self.must_open_db)
760a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
770a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # Get "to be recovered" transactions but
780a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # let them be garbage collected.
790a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        recovered_txns=self.dbenv.txn_recover()
800a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.assertEqual(self.num_txns,len(recovered_txns))
810a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for gid,txn in recovered_txns :
820a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.assertTrue(gid in txns)
830a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        del txn
840a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        del recovered_txns
850a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
860a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self._recreate_env(self.must_open_db)
870a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
880a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # Get "to be recovered" transactions. Commit, abort and
890a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # discard them.
900a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        recovered_txns=self.dbenv.txn_recover()
910a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.assertEqual(self.num_txns,len(recovered_txns))
920a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        discard_txns=set()
930a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        committed_txns=set()
940a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        state=0
950a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for gid,txn in recovered_txns :
960a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            if state==0 or state==1:
970a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                committed_txns.add(gid)
980a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                txn.commit()
990a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            elif state==2 :
1000a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                txn.abort()
1010a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            elif state==3 :
1020a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                txn.discard()
1030a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                discard_txns.add(gid)
1040a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                state=-1
1050a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            state+=1
1060a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        del txn
1070a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        del recovered_txns
1080a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1090a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self._recreate_env(self.must_open_db)
1100a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1110a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # Verify the discarded transactions are still
1120a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # around, and dispose them.
1130a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        recovered_txns=self.dbenv.txn_recover()
1140a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.assertEqual(len(discard_txns),len(recovered_txns))
1150a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for gid,txn in recovered_txns :
1160a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            txn.abort()
1170a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        del txn
1180a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        del recovered_txns
1190a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1200a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self._recreate_env(must_open_db=True)
1210a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1220a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # Be sure there are not pending transactions.
1230a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # Check also database size.
1240a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        recovered_txns=self.dbenv.txn_recover()
1250a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.assertTrue(len(recovered_txns)==0)
1260a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.assertEqual(len(committed_txns),self.db.stat()["nkeys"])
1270a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1280a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass DBTxn_distributedSYNC(DBTxn_distributed):
1290a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    nosync=False
1300a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1310a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass DBTxn_distributed_must_open_db(DBTxn_distributed):
1320a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    must_open_db=True
1330a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1340a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass DBTxn_distributedSYNC_must_open_db(DBTxn_distributed):
1350a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    nosync=False
1360a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    must_open_db=True
1370a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1380a8c90248264a8b26970b4473770bcc3df8515fJosh Gao#----------------------------------------------------------------------
1390a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1400a8c90248264a8b26970b4473770bcc3df8515fJosh Gaodef test_suite():
1410a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    suite = unittest.TestSuite()
1420a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    if db.version() >= (4,5) :
1430a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        suite.addTest(unittest.makeSuite(DBTxn_distributed))
1440a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        suite.addTest(unittest.makeSuite(DBTxn_distributedSYNC))
1450a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    if db.version() >= (4,6) :
1460a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        suite.addTest(unittest.makeSuite(DBTxn_distributed_must_open_db))
1470a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        suite.addTest(unittest.makeSuite(DBTxn_distributedSYNC_must_open_db))
1480a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    return suite
1490a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1500a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1510a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoif __name__ == '__main__':
1520a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    unittest.main(defaultTest='test_suite')
153