Python recordio.ThriftRecordWriter类代码示例

def test_basic_thriftrecordwriter_write():
  test_string = StringType("hello world")

  with EphemeralFile('w') as fp:
    fn = fp.name

    rw = ThriftRecordWriter(fp)
    rw.write(test_string)
    rw.close()

    with open(fn) as fpr:
      rr = ThriftRecordReader(fpr, StringType)
      assert rr.read() == test_string

def test_thrift_recordwriter_type_mismatch():
  test_string = StringType("hello world")
  with EphemeralFile('w') as fp:
    fn = fp.name

    rw = ThriftRecordWriter(fp)
    rw.write(test_string)
    rw.close()

    with open(fn) as fpr:
      rr = ThriftRecordReader(fpr, IntType)
      # This is a peculiar behavior of Thrift in that it just returns
      # ThriftType() with no serialization applied
      assert rr.read() == IntType()

def test_paranoid_thrift_append_framing():
  test_string_1 = StringType("hello world")
  test_string_2 = StringType("ahoy ahoy, bonjour")

  with EphemeralFile('w') as fp:
    fn = fp.name

    ThriftRecordWriter.append(fn, test_string_1)
    ThriftRecordWriter.append(fn, test_string_2)

    with open(fn) as fpr:
      rr = ThriftRecordReader(fpr, StringType)
      assert rr.read() == test_string_1
      assert rr.read() == test_string_2

def test_thriftrecordreader_iteration():
  test_string_1 = StringType("hello world")
  test_string_2 = StringType("ahoy ahoy, bonjour")

  with EphemeralFile('w') as fp:
    fn = fp.name

    rw = ThriftRecordWriter(fp)
    rw.write(test_string_1)
    rw.write(test_string_2)
    rw.close()

    with open(fn) as fpr:
      rr = ThriftRecordReader(fpr, StringType)
      records = []
      for record in rr:
        records.append(record)
      assert records == [test_string_1, test_string_2]

 def _setup_ckpt(self):
   """Set up the checkpoint: must be run on the parent."""
   self._log('initializing checkpoint file: %s' % self.ckpt_file())
   ckpt_fp = lock_file(self.ckpt_file(), "a+")
   if ckpt_fp in (None, False):
     raise self.CheckpointError('Could not acquire checkpoint permission or lock for %s!' %
       self.ckpt_file())
   self._ckpt_head = os.path.getsize(self.ckpt_file())
   ckpt_fp.seek(self._ckpt_head)
   self._ckpt = ThriftRecordWriter(ckpt_fp)
   self._ckpt.set_sync(True)

 def open_checkpoint(cls, filename, force=False, state=None):
     """
   Acquire a locked checkpoint stream.
 """
     safe_mkdir(os.path.dirname(filename))
     fp = lock_file(filename, "a+")
     if fp in (None, False):
         if force:
             log.info("Found existing runner, forcing leadership forfeit.")
             state = state or CheckpointDispatcher.from_file(filename)
             if cls.kill_runner(state):
                 log.info("Successfully killed leader.")
                 # TODO(wickman)  Blocking may not be the best idea here.  Perhaps block up to
                 # a maximum timeout.  But blocking is necessary because os.kill does not immediately
                 # release the lock if we're in force mode.
                 fp = lock_file(filename, "a+", blocking=True)
         else:
             log.error("Found existing runner, cannot take control.")
     if fp in (None, False):
         raise cls.PermissionError("Could not open locked checkpoint: %s, lock_file = %s" % (filename, fp))
     ckpt = ThriftRecordWriter(fp)
     ckpt.set_sync(True)
     return ckpt

def test_thriftrecordwriter_framing():
  test_string_1 = StringType("hello world")
  test_string_2 = StringType("ahoy ahoy, bonjour")

  with EphemeralFile('w') as fp:
    fn = fp.name

    rw = ThriftRecordWriter(fp)
    rw.write(test_string_1)
    rw.close()

    with open(fn, 'a') as fpa:
      rw = ThriftRecordWriter(fpa)
      rw.write(test_string_2)

    with open(fn) as fpr:
      rr = ThriftRecordReader(fpr, StringType)
      assert rr.read() == test_string_1
      assert rr.read() == test_string_2

#.........这里部分代码省略.........
                               fork_time=self._fork_time,
                               coordinator_pid=self._pid)

  def cmdline(self):
    return self._cmdline

  def name(self):
    return self._name

  def pid(self):
    """pid of the coordinator"""
    return self._pid

  def rebind(self, pid, fork_time):
    """rebind Process to an existing coordinator pid without forking"""
    self._pid = pid
    self._fork_time = fork_time

  def ckpt_file(self):
    return self._pathspec.getpath('process_checkpoint')

  def process_logdir(self):
    return self._pathspec.getpath('process_logdir')

  def _setup_ckpt(self):
    """Set up the checkpoint: must be run on the parent."""
    self._log('initializing checkpoint file: %s' % self.ckpt_file())
    ckpt_fp = lock_file(self.ckpt_file(), "a+")
    if ckpt_fp in (None, False):
      raise self.CheckpointError('Could not acquire checkpoint permission or lock for %s!' %
        self.ckpt_file())
    self._ckpt_head = os.path.getsize(self.ckpt_file())
    ckpt_fp.seek(self._ckpt_head)
    self._ckpt = ThriftRecordWriter(ckpt_fp)
    self._ckpt.set_sync(True)

  def _init_ckpt_if_necessary(self):
    if self._ckpt is None:
      self._setup_ckpt()

  def _wait_for_control(self):
    """Wait for control of the checkpoint stream: must be run in the child."""
    total_wait_time = Amount(0, Time.SECONDS)

    with open(self.ckpt_file(), 'r') as fp:
      fp.seek(self._ckpt_head)
      rr = ThriftRecordReader(fp, RunnerCkpt)
      while total_wait_time < self.MAXIMUM_CONTROL_WAIT:
        ckpt_tail = os.path.getsize(self.ckpt_file())
        if ckpt_tail == self._ckpt_head:
          self._platform.clock().sleep(self.CONTROL_WAIT_CHECK_INTERVAL.as_(Time.SECONDS))
          total_wait_time += self.CONTROL_WAIT_CHECK_INTERVAL
          continue
        checkpoint = rr.try_read()
        if checkpoint:
          if not checkpoint.process_status:
            raise self.CheckpointError('No process status in checkpoint!')
          if (checkpoint.process_status.process != self.name() or
              checkpoint.process_status.state != ProcessState.FORKED or
              checkpoint.process_status.fork_time != self._fork_time or
              checkpoint.process_status.coordinator_pid != self._pid):
            self._log('Losing control of the checkpoint stream:')
            self._log('   fork_time [%s] vs self._fork_time [%s]' % (
                checkpoint.process_status.fork_time, self._fork_time))
            self._log('   coordinator_pid [%s] vs self._pid [%s]' % (
                checkpoint.process_status.coordinator_pid, self._pid))