diff --git a/centipede/centipede_callbacks.cc b/centipede/centipede_callbacks.cc index 25293c78..4b6ba0af 100644 --- a/centipede/centipede_callbacks.cc +++ b/centipede/centipede_callbacks.cc @@ -215,7 +215,8 @@ int CentipedeCallbacks::ExecuteCentipedeSancovBinaryWithShmem( // Get results. batch_result.exit_code() = retval; - CHECK(batch_result.Read(outputs_blobseq_)); + const bool read_success = batch_result.Read(outputs_blobseq_); + LOG_IF(ERROR, !read_success) << "Failed to read batch result!"; outputs_blobseq_.ReleaseSharedMemory(); // Outputs are already consumed. // We may have fewer feature blobs than inputs if @@ -225,7 +226,8 @@ int CentipedeCallbacks::ExecuteCentipedeSancovBinaryWithShmem( // * Will be logged by the caller. // * some outputs were not written because the outputs_blobseq_ overflown. // * Logged by the following code. - if (retval == 0 && batch_result.num_outputs_read() != num_inputs_written) { + if (retval == 0 && read_success && + batch_result.num_outputs_read() != num_inputs_written) { LOG(INFO) << "Read " << batch_result.num_outputs_read() << "/" << num_inputs_written << " outputs; shmem_size_mb might be too small: " diff --git a/centipede/runner_result.cc b/centipede/runner_result.cc index 0dd27a51..4c9bdb9d 100644 --- a/centipede/runner_result.cc +++ b/centipede/runner_result.cc @@ -101,15 +101,18 @@ bool BatchResult::Read(BlobSequence &blobseq) { continue; } if (blob.tag == kTagMetadata) { + if (current_execution_result == nullptr) return false; current_execution_result->metadata().Read(blob); continue; } if (blob.tag == kTagStats) { + if (current_execution_result == nullptr) return false; if (blob.size != sizeof(ExecutionResult::Stats)) return false; memcpy(¤t_execution_result->stats(), blob.data, blob.size); continue; } if (blob.tag == kTagFeatures) { + if (current_execution_result == nullptr) return false; const size_t features_size = blob.size / sizeof(feature_t); FeatureVec &features = current_execution_result->mutable_features(); features.resize(features_size); @@ -117,7 +120,6 @@ bool BatchResult::Read(BlobSequence &blobseq) { features_size * sizeof(feature_t)); } } - num_outputs_read_ = num_ends; return true; } diff --git a/centipede/runner_result_test.cc b/centipede/runner_result_test.cc index dc6d037f..f748e7fc 100644 --- a/centipede/runner_result_test.cc +++ b/centipede/runner_result_test.cc @@ -194,5 +194,26 @@ TEST(MutationResult, WriteThenRead) { ElementsAre(ByteArray{1, 2, 3}, ByteArray{4, 5, 6}, ByteArray{7, 8, 9})); } +TEST(ExecutionResult, ReadResultSucceedsOnlyWithInputBegin) { + auto buffer = std::make_unique(1000); + BlobSequence blobseq(buffer.get(), 1000); + BatchResult batch_result; + + EXPECT_TRUE(BatchResult::WriteInputBegin(blobseq)); + EXPECT_TRUE(BatchResult::WriteOneFeatureVec({}, 0, blobseq)); + EXPECT_TRUE(BatchResult::WriteInputEnd(blobseq)); + blobseq.Reset(); + batch_result.ClearAndResize(1); + EXPECT_TRUE(batch_result.Read(blobseq)); + + blobseq.Reset(); + EXPECT_TRUE(BatchResult::WriteOneFeatureVec({}, 0, blobseq)); + EXPECT_TRUE(BatchResult::WriteInputEnd(blobseq)); + + blobseq.Reset(); + batch_result.ClearAndResize(1); + EXPECT_FALSE(batch_result.Read(blobseq)); +} + } // namespace } // namespace centipede