Make buffering consuemr thread-safe. It's ugly, but it works for now. We'll switch to something else.

Author: pmenon

src/codegen/buffering_consumer.cpp

@@ -57,27 +57,30 @@ BufferingConsumer::BufferingConsumer(const std::vector<oid_t> &cols,
   for (oid_t col_id : cols) {
     output_ais_.push_back(context.Find(col_id));
   }
-  state.output = &tuples_;
 }
 
-// Append the array of values (i.e., a tuple) into the consumer's buffer of
-// output tuples.
-void BufferingConsumer::BufferTuple(char *state, char *tuple,
+// Append the array of row attributes to the buffer of tuples
+//
+// Note: buffering consumers rely on an ugly mutex to protect access to the
+//       output buffer. This is ugly AF. We don't actually use it for primary
+//       query processing, so it's okay.
+void BufferingConsumer::BufferTuple(char *opaque_state, char *tuple,
                                     uint32_t num_cols) {
-  BufferingState *buffer_state = reinterpret_cast<BufferingState *>(state);
-  buffer_state->output->emplace_back(
-      reinterpret_cast<peloton::type::Value *>(tuple), num_cols);
+  auto *buffer = reinterpret_cast<Buffer *>(opaque_state);
+  std::lock_guard<std::mutex> lock{buffer->mutex};
+  buffer->output.emplace_back(reinterpret_cast<peloton::type::Value *>(tuple),
+                              num_cols);
 }
 
 // Create two pieces of state: a pointer to the output tuple vector and an
 // on-stack value array representing a single tuple.
-void BufferingConsumer::Prepare(CompilationContext &ctx) {
+void BufferingConsumer::Prepare(CompilationContext &compilation_ctx) {
   // Be sure to call our parent
-  ExecutionConsumer::Prepare(ctx);
+  ExecutionConsumer::Prepare(compilation_ctx);
 
   // Install a little char* for the state we need
-  CodeGen &codegen = ctx.GetCodeGen();
-  QueryState &query_state = ctx.GetQueryState();
+  CodeGen &codegen = compilation_ctx.GetCodeGen();
+  QueryState &query_state = compilation_ctx.GetQueryState();
   consumer_state_id_ =
       query_state.RegisterState("consumerState", codegen.CharPtrType());
 }
@@ -137,12 +140,23 @@ void BufferingConsumer::ConsumeResult(ConsumerContext &ctx,
     codegen.CallFunc(output_func, args);
   }
 
+  auto &query_state = ctx.GetQueryState();
+  llvm::Value *buffer_ptr =
+      query_state.LoadStateValue(codegen, consumer_state_id_);
+
   // Append the tuple to the output buffer (by calling BufferTuple(...))
-  auto *consumer_state = GetStateValue(ctx, consumer_state_id_);
-  std::vector<llvm::Value *> args = {consumer_state, tuple_buffer_,
+  std::vector<llvm::Value *> args = {buffer_ptr, tuple_buffer_,
                                      codegen.Const32(output_ais_.size())};
   codegen.Call(BufferingConsumerProxy::BufferTuple, args);
 }
 
+char *BufferingConsumer::GetConsumerState() {
+  return reinterpret_cast<char *>(&buffer_);
+}
+
+const std::vector<WrappedTuple> &BufferingConsumer::GetOutputTuples() const {
+  return buffer_.output;
+}
+
 }  // namespace codegen
 }  // namespace peloton

src/include/codegen/buffering_consumer.h

@@ -6,13 +6,14 @@
 //
 // Identification: src/include/codegen/buffering_consumer.h
 //
-// Copyright (c) 2015-2017, Carnegie Mellon University Database Group
+// Copyright (c) 2015-2018, Carnegie Mellon University Database Group
 //
 //===----------------------------------------------------------------------===//
 
 #pragma once
 
 #include <vector>
+#include <mutex>
 
 #include "codegen/compilation_context.h"
 #include "codegen/execution_consumer.h"
@@ -50,51 +51,41 @@ class WrappedTuple : public ContainerTuple<std::vector<peloton::type::Value>> {
 //===----------------------------------------------------------------------===//
 class BufferingConsumer : public ExecutionConsumer {
  public:
-  struct BufferingState {
-    std::vector<WrappedTuple> *output;
-  };
-
-  // Constructor
+  /// Constructor
   BufferingConsumer(const std::vector<oid_t> &cols,
                     const planner::BindingContext &context);
 
-  void Prepare(CompilationContext &compilation_context) override;
+  void Prepare(CompilationContext &compilation_ctx) override;
 
   // Query state
   void InitializeQueryState(CompilationContext &) override {}
   void TearDownQueryState(CompilationContext &) override {}
 
-  // TODO(pmenon): Implement me
-  bool SupportsParallelExec() const override { return false; }
+  bool SupportsParallelExec() const override { return true; }
 
   void ConsumeResult(ConsumerContext &ctx, RowBatch::Row &row) const override;
 
-  llvm::Value *GetStateValue(ConsumerContext &ctx,
-                             const QueryState::Id &id) const {
-    auto &query_state = ctx.GetQueryState();
-    return query_state.LoadStateValue(ctx.GetCodeGen(), id);
-  }
-
   // Called from compiled query code to buffer the tuple
-  static void BufferTuple(char *state, char *tuple, uint32_t num_cols);
+  static void BufferTuple(char *buffer, char *tuple, uint32_t num_cols);
 
   //===--------------------------------------------------------------------===//
   // ACCESSORS
   //===--------------------------------------------------------------------===//
 
-  char *GetConsumerState() override { return reinterpret_cast<char *>(&state); }
+  char *GetConsumerState() override;
 
-  const std::vector<WrappedTuple> &GetOutputTuples() const { return tuples_; }
+  const std::vector<WrappedTuple> &GetOutputTuples() const;
 
  private:
   // The attributes we want to output
   std::vector<const planner::AttributeInfo *> output_ais_;
 
-  // Buffered output tuples
-  std::vector<WrappedTuple> tuples_;
-
-  // Running buffering state
-  BufferingState state;
+  // The thread-safe buffer of output tuples
+  struct Buffer {
+    std::mutex mutex;
+    std::vector<WrappedTuple> output;
+  };
+  Buffer buffer_;
 
   // The slot in the runtime state to find our state context
   QueryState::Id consumer_state_id_;