简单线程池

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <assert.h>


// 任务结构体，使用队列保存
typedef struct worker_s worker_t;
struct worker_s
{
    void *(*process)(void *arg);
    void *arg;
    worker_t *next;
};

// 线程池结构
typedef struct pool_s pool_t;
struct pool_s
{
    pthread_mutex_t mutex;
    pthread_cond_t cond;

    worker_t *head;
    worker_t *tail;

    pthread_t *ids;
    int max_thread_num;
};


void *thread_cb(void *arg)
{
    pool_t *pool = (pool_t*)arg;
    long self = pthread_self();

    while (1)
    {
        int ret;
        if ((ret = pthread_mutex_lock(&pool->mutex)) != 0)
        {
            fprintf(stderr, "pthread_mutex_lock: %d
", ret);
            break;
        }

        printf("in %ld
", self);
        while (pool->head == NULL)
        {
            printf("wait %ld
", self);
            if ((ret = pthread_cond_wait(&pool->cond, &pool->mutex)) != 0)
            {
                fprintf(stderr, "pthread_cond_wait: %d
", ret);
                pthread_mutex_unlock(&pool->mutex);
            }
        }

        assert(pool->head != NULL);

        // 取出队列保存的第一个任务
        worker_t *w = pool->head;
        pool->head = pool->head->next;
        pthread_mutex_unlock(&pool->mutex);

        w->process(w->arg);
        free(w);
    }

    return NULL;
}

pool_t *pool_init(int thread_num)
{
    pool_t *pool = (pool_t*)malloc(sizeof(pool_t));
    if (pool)
    {
        if (pthread_mutex_init(&pool->mutex, NULL) != 0)
        {
            return NULL;
        }
        if (pthread_cond_init(&pool->cond, NULL) != 0)
        {
            pthread_mutex_destroy(&pool->mutex);
            return NULL;
        }
        pool->head = pool->tail = NULL;
        pool->max_thread_num = thread_num;
        pool->ids = (pthread_t*)malloc(sizeof(pthread_t) * thread_num);

        for (int i = 0; i < thread_num; i++)
        {
            pthread_create(&pool->ids[i], NULL, thread_cb, (void*)pool);
            pthread_detach(pool->ids[i]);
        }
    }

    return pool;
}

int pool_add_worker(pool_t *pool, void *(*process)(void *arg), void *arg)
{
    worker_t *w = (worker_t*)malloc(sizeof(worker_t));
    if (NULL == w)
    {
        return -1;
    }
    w->process = process;
    w->arg = arg;
    w->next = NULL;

    pthread_mutex_lock(&pool->mutex);
    if (pool->head != NULL)
    {
        pool->tail->next = w;
        pool->tail = w;
    }
    else
    {
        pool->head = pool->tail = w;
    }

    pthread_cond_signal(&pool->cond);
    pthread_mutex_unlock(&pool->mutex);

    return 0;
}


void *p_cb(void *arg)
{
    sleep(2);
    printf("callback func print
");

    return NULL;
}


int main()
{
    pool_t * pool = pool_init(3);

    for (int i = 0; i < 5; ++i)
    {
        if (pool_add_worker(pool, p_cb, NULL) != 0)
        {
            fprintf(stderr, "add worker error
");
        }
    }

    getchar();

    return 0;
}

 C++版本的：

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <functional>
#include <future>
#include <iostream>
#include <mutex>
#include <queue>
#include <thread>
#include <vector>

class Pool {
 public:
  Pool(int n = std::thread::hardware_concurrency()) {
    for (int i = 0; i < n; ++i) {
      threads_.emplace_back(&Pool::ThreadCb, this);
    }
  }
  ~Pool() { Stop(); }

  // 添加任务
  template <class F, class... Args>
  void AddTask(F&& f, Args&&... args) {
    {
      std::lock_guard<std::mutex> lk(in_fly_mutex_);
      ++in_fly_;
    }

    std::lock_guard<std::mutex> lk(task_queue_mutex_);
    task_queue_.emplace(
        std::bind(std::forward<F>(f), std::forward<Args>(args)...));
    task_queue_cv_.notify_one();
  }

  // 等待线程池中的任务执行结束
  void Wait() {
    std::unique_lock<std::mutex> lk(in_fly_mutex_);
    in_fly_cv_.wait(lk, [this] { return this->in_fly_ == 0; });
    printf("-- %d --
", in_fly_);
  }

 private:
  // 执行完队列中所有的任务后，再退出
  void Stop() {
    stop_ = true;
    {
      std::lock_guard<std::mutex> lk(task_queue_mutex_);
      task_queue_cv_.notify_all();
    }

    for (auto& e : threads_) {
      e.join();
    }

    threads_.clear();
  }

  void ThreadCb() {
    while (1) {
      std::unique_lock<std::mutex> lk(task_queue_mutex_);
      task_queue_cv_.wait(lk, [this] { return !task_queue_.empty() || stop_; });
      // 只要队列中还有任务，就不退出
      if (!task_queue_.empty()) {
        auto task = task_queue_.front();
        task_queue_.pop();
        lk.unlock();

        InFlyGuard(*this);
        task();
      } else {
        return;
      }
    }
  }

 private:
  bool stop_ = false;
  std::queue<std::function<void(void)>> task_queue_;
  std::mutex task_queue_mutex_;
  std::condition_variable task_queue_cv_;
  std::vector<std::thread> threads_;

 private:
  struct InFlyGuard {
    InFlyGuard(Pool& pool) : pool_(pool) {}
    ~InFlyGuard() {
      std::lock_guard<std::mutex> lk(pool_.in_fly_mutex_);
      --pool_.in_fly_;
      if (pool_.in_fly_ == 0) {
        pool_.in_fly_cv_.notify_all();
      }
    }
    Pool& pool_;
  };

  // 尚未结束的任务数
  int in_fly_ = 0;
  std::mutex in_fly_mutex_;
  std::condition_variable in_fly_cv_;
};

std::atomic<int> count{0};

void Add(int a, int b) {
  ++count;
  printf("%d
", count.load(std::memory_order::memory_order_relaxed));
  std::this_thread::sleep_for(std::chrono::milliseconds(800));
}

void Sub(int a, int b) {
  --count;
  printf("%d
", count.load(std::memory_order::memory_order_relaxed));
  std::this_thread::sleep_for(std::chrono::milliseconds(800));
}

int main() {
  Pool p(3);

  p.AddTask(Add, 2, 3);
  p.AddTask(Add, 2, 3);
  p.AddTask(Add, 2, 3);
  p.AddTask(Add, 2, 3);
  p.AddTask(Add, 2, 3);

  // 等待上面的任务结束
  p.Wait();
  printf("main thread, count: %d
",
         count.load(std::memory_order::memory_order_relaxed));

  p.AddTask(Sub, 9, 7);
  p.AddTask(Sub, 9, 7);
  p.AddTask(Sub, 9, 7);
  p.AddTask(Sub, 9, 7);
  p.AddTask(Sub, 9, 7);

  return 0;
}

 这段代码的 Wait 函数不符合预期，对比下面的代码：

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <functional>
#include <future>
#include <iostream>
#include <memory>
#include <mutex>
#include <queue>
#include <thread>
#include <vector>

class Pool {
 public:
  Pool(int n = std::thread::hardware_concurrency()) : size_(n) {
    for (int i = 0; i < n; ++i) {
      pool_.emplace_back(&Pool::thread_func, this);
    }
  }

  ~Pool() { Stop(); }

  template <class Func, class... Args>
  void AddTask(Func &&f, Args &&... args) {
    {
      std::unique_lock<std::mutex> lk(in_flight_mutex_);
      ++in_flight_;
    }
    std::unique_lock<std::mutex> lk(task_queue_mutex_);
    task_queue_.emplace(
        std::bind(std::forward<Func>(f), std::forward<Args>(args)...));
    task_queue_cv_.notify_one();
  }

  void Wait() {
    std::unique_lock<std::mutex> lk(in_flight_mutex_);
    in_flight_cv_.wait(lk, [this] { return in_flight_ == 0; });
  }

 private:
  struct FlyGuard {
    FlyGuard(Pool &p) : p_(p) {}
    ~FlyGuard() {
      std::unique_lock<std::mutex> lk(p_.in_flight_mutex_);
      --p_.in_flight_;
      if (p_.in_flight_ == 0) {
        p_.in_flight_cv_.notify_all();
      }
    }

    Pool &p_;
  };
  void thread_func() {
    while (1) {
      std::unique_lock<std::mutex> lk(task_queue_mutex_);
      task_queue_cv_.wait(lk, [this] { return stop_ || !task_queue_.empty(); });
      if (!task_queue_.empty()) {
        auto task = task_queue_.front();
        task_queue_.pop();
        lk.unlock();

        FlyGuard fg(*this);
        task();
      } else {
        return;
      }
    }
  }
  void Stop() {
    puts("stop");
    stop_ = true;
    std::unique_lock<std::mutex> lk(task_queue_mutex_);
    task_queue_cv_.notify_all();
    lk.unlock();

    for (auto &e : pool_) {
      e.join();
    }
    pool_.clear();
  }

 private:
  bool stop_ = false;
  std::vector<std::thread> pool_;
  int size_;
  std::mutex task_queue_mutex_;
  std::condition_variable task_queue_cv_;
  std::queue<std::function<void(void)>> task_queue_;

  std::mutex in_flight_mutex_;
  std::condition_variable in_flight_cv_;
  int in_flight_ = 0;
};

std::atomic<int> count{0};

void Add(int a, int b) {
  ++count;
  printf("%d
", count.load(std::memory_order::memory_order_relaxed));
  std::this_thread::sleep_for(std::chrono::milliseconds(800));
}

void Sub(int a, int b) {
  --count;
  printf("%d
", count.load(std::memory_order::memory_order_relaxed));
  std::this_thread::sleep_for(std::chrono::milliseconds(800));
}

int main() {
  Pool p(3);

  p.AddTask(Add, 2, 3);
  p.AddTask(Add, 2, 3);
  p.AddTask(Add, 2, 3);
  p.AddTask(Add, 2, 3);
  p.AddTask(Add, 2, 3);

  // 等待上面的任务结束
  p.Wait();
  printf("main thread, count: %d
",
         count.load(std::memory_order::memory_order_relaxed));

  p.AddTask(Sub, 9, 7);
  p.AddTask(Sub, 9, 7);
  p.AddTask(Sub, 9, 7);
  p.AddTask(Sub, 9, 7);
  p.AddTask(Sub, 9, 7);

  return 0;
}

这一段就是在Wait那里明显停顿了一下，符合预期。

这两者代码逻辑一样，为啥表现不一样。。不得其解

调试了半天，发现是这里写错了

InFlyGuard(*this);