dbms: development [#CONV-2944].

dbms/include/DB/DataStreams/AsynchronousBlockInputStream.h

@@ -1,8 +1,8 @@
 #pragma once
 
-#include <Poco/Mutex.h>
-#include <Poco/Thread.h>
-#include <Poco/Runnable.h>
+#include <statdaemons/threadpool.hpp>
+
+#include <Poco/Event.h>
 
 #include <DB/DataStreams/IProfilingBlockInputStream.h>
 
@@ -10,53 +10,6 @@
 namespace DB
 {
 
-/** thread-safe очередь из одного элемента,
-  * рассчитанная на одного producer-а и одного consumer-а.
-  */
-template <typename T>
-class OneElementQueue
-{
-private:
-	T data;
-	Poco::FastMutex mutex_fill;		/// Захвачен, когда данные есть.
-	Poco::FastMutex mutex_empty;	/// Захвачен, когда данных нет.
-
-public:
-	OneElementQueue()
-	{
-		mutex_empty.lock();
-	}
-
-	/// Вызывается единственным producer-ом.
-	void push(const T & x)
-	{
-		mutex_fill.lock();
- 		data = x;
-		mutex_empty.unlock();
-	}
-
-	/// Вызывается единственным consumer-ом.
-	void pop(T & x)
-	{
-		mutex_empty.lock();
-		x = data;
-		mutex_fill.unlock();
-	}
-
-	/// Позволяет ждать элемента не дольше заданного таймаута. Вызывается единственным consumer-ом.
-	bool poll(UInt64 milliseconds)
-	{
-		if (mutex_empty.tryLock(milliseconds))
-		{
-			mutex_empty.unlock();
-			return true;
-		}
-
-		return false;
-	}
-};
-
-
 /** Выполняет другой BlockInputStream в отдельном потоке.
   * Это служит для двух целей:
   * 1. Позволяет сделать так, чтобы разные стадии конвеьера выполнения запроса работали параллельно.
@@ -67,74 +20,87 @@ public:
 class AsynchronousBlockInputStream : public IProfilingBlockInputStream
 {
 public:
-	AsynchronousBlockInputStream(BlockInputStreamPtr in_) : in(in_), started(false), runnable(*this)
+	AsynchronousBlockInputStream(BlockInputStreamPtr in_) : in(in_), pool(1), started(false)
 	{
 		children.push_back(in);
 	}
 
-
-	/** Ждать готовность данных не более заданного таймаута. Запустить получение данных, если нужно.
-	  * Если функция вернула true - данные готовы и можно делать read().
-	  */
-	bool poll(UInt64 milliseconds)
-	{
-		startIfNeed();
-		return output_queue.poll(milliseconds);
-	}
-	
-
 	String getName() const { return "AsynchronousBlockInputStream"; }
 
 	BlockInputStreamPtr clone() { return new AsynchronousBlockInputStream(in); }
 
 
-    ~AsynchronousBlockInputStream()
-	{
-		if (started)
-			thread->join();
-	}
-
-protected:
-	Block readImpl()
-	{
-		OutputData res;
-
-		startIfNeed();
-
-		/// Будем ждать, пока будет готов следующий блок или будет выкинуто исключение.
-		output_queue.pop(res);
-
-		if (res.exception)
-			res.exception->rethrow();
-
-		return res.block;
-	}
-	
-	
-	void startIfNeed()
+	/** Ждать готовность данных не более заданного таймаута. Запустить получение данных, если нужно.
+	  * Если функция вернула true - данные готовы и можно делать read(); нельзя вызвать функцию сразу ещё раз.
+	  */
+	bool poll(UInt64 milliseconds)
 	{
 		if (!started)
 		{
+			next();
 			started = true;
-			thread = new Poco::Thread;
-			thread->start(runnable);
 		}
+
+		return ready.tryWait(milliseconds);
+	}
+
+
+    ~AsynchronousBlockInputStream()
+	{
+		if (started)
+			pool.wait();
+	}
+
+protected:
+	BlockInputStreamPtr in;
+	boost::threadpool::pool pool;
+	Poco::Event ready;
+	bool started;
+
+	Block block;
+	ExceptionPtr exception;
+
+	
+	Block readImpl()
+	{
+		/// Если вычислений ещё не было - вычислим первый блок синхронно
+		if (!started)
+		{
+			calculate();
+			started = true;
+		}
+		else	/// Если вычисления уже идут - подождём результата
+			pool.wait();
+
+		if (exception)
+			exception->rethrow();
+
+		Block res = block;
+		if (!res)
+			return res;
+
+		/// Запустим вычисления следующего блока
+		block = Block();
+		next();
+
+		return res;
+	}
+	
+
+	void next()
+	{
+		pool.schedule(boost::bind(&AsynchronousBlockInputStream::calculate, this));
 	}
 	
 
 	/// Вычисления, которые могут выполняться в отдельном потоке
-	class Thread : public Poco::Runnable
+	void calculate()
 	{
-	public:
-		Thread(AsynchronousBlockInputStream & parent_) : parent(parent_) {}
-
-		void run()
-		{
-			ExceptionPtr exception;
+		ready.reset();
 
 		try
 		{
-				loop();
+			block = in->read();
 		}
 		catch (const Exception & e)
 		{
@@ -153,49 +119,8 @@ protected:
 			exception = new Exception("Unknown exception", ErrorCodes::UNKNOWN_EXCEPTION);
 		}
 
-			if (exception)
-			{
-				parent.cancel();
-
-				/// Отдаём эксепшен в основной поток.
-				parent.output_queue.push(exception);
+		ready.set();
 	}
-		}
-
-		void loop()
-		{
-			while (true)
-			{
-				Block res = parent.in->read();
-				parent.output_queue.push(res);
-
-				if (!res)
-					break;
-			}
-		}
-
-	private:
-		AsynchronousBlockInputStream & parent;
-	};
-
-
-	BlockInputStreamPtr in;
-	bool started;
-
-	struct OutputData
-	{
-		Block block;
-		ExceptionPtr exception;
-
-		OutputData() {}
-		OutputData(Block & block_) : block(block_) {}
-		OutputData(ExceptionPtr & exception_) : exception(exception_) {}
-	};
-
-	OneElementQueue<OutputData> output_queue;
-
-	Thread runnable;
-	SharedPtr<Poco::Thread> thread;
 };
 
 }