Peiyu/fix engine unpickling bug

src/python/pants/engine/exp/engine.py

@@ -5,6 +5,7 @@
 from __future__ import (absolute_import, division, generators, nested_scopes, print_function,
                         unicode_literals, with_statement)
 
+import cPickle as pickle
 import functools
 import logging
 import multiprocessing
@@ -21,12 +22,6 @@
 from pants.util.objects import datatype
 
 
-try:
-  import cPickle as pickle
-except ImportError:
-  import pickle
-
-
 logger = logging.getLogger(__name__)
 
 

src/python/pants/engine/exp/storage.py

@@ -6,14 +6,15 @@
                         unicode_literals, with_statement)
 
 import cPickle as pickle
-import cStringIO as StringIO
 import sys
 from abc import abstractmethod
 from binascii import hexlify
 from collections import Counter
 from contextlib import closing
 from functools import total_ordering
 from hashlib import sha1
+from io import BytesIO
+from StringIO import StringIO
 from struct import Struct as StdlibStruct
 
 import lmdb
@@ -166,7 +167,7 @@ def put(self, obj):
     Longer term see https://github.com/pantsbuild/pants/issues/2969
     """
     try:
-      with closing(StringIO.StringIO()) as buf:
+      with closing(BytesIO()) as buf:
         pickler = pickle.Pickler(buf, protocol=self._protocol)
         pickler.fast = 1
         pickler.dump(obj)
@@ -493,7 +494,7 @@ def get(self, key):
     with self._env.begin(db=self._db, buffers=True) as txn:
       value = txn.get(key)
       if value is not None:
-        return StringIO.StringIO(value)
+        return StringIO(value)
       return None
 
   def put(self, key, value):

tests/python/pants_test/engine/exp/test_engine.py

@@ -46,7 +46,6 @@ def test_serial_engine_simple(self):
     engine = LocalSerialEngine(self.scheduler, self.storage, self.cache)
     self.assert_engine(engine)
 
-  @unittest.skip('https://github.com/pantsbuild/pants/issues/3149')
   def test_multiprocess_engine_multi(self):
     with self.multiprocessing_engine() as engine:
       self.assert_engine(engine)
@@ -62,7 +61,6 @@ def test_multiprocess_unpickleable(self):
       with self.assertRaises(SerializationError):
         engine.execute(build_request)
 
-  @unittest.skip('https://github.com/pantsbuild/pants/issues/3149')
   def test_rerun_with_cache(self):
     with self.multiprocessing_engine() as engine:
       self.assert_engine(engine)