45 lines
884 B
C++
45 lines
884 B
C++
#pragma once
|
|
|
|
#include <vector>
|
|
#include <atomic>
|
|
#include <thread>
|
|
|
|
#include "shm_helper.h"
|
|
|
|
|
|
class ShmWorker {
|
|
public:
|
|
ShmWorker();
|
|
~ShmWorker();
|
|
|
|
bool register_worker(int32_t tgid, int gpu_id, uint64_t *out_shareable_handle,
|
|
uint64_t *out_vmem_size);
|
|
|
|
bool try_lock_gpu(bool &out_self_hold);
|
|
bool lock_gpu(bool &out_self_hold);
|
|
void unlock_gpu(bool keep_wait = false);
|
|
|
|
bool has_higher_priority_waiter();
|
|
void start_wait();
|
|
void cancel_wait();
|
|
|
|
private:
|
|
int32_t tgid;
|
|
int gpu_id;
|
|
int shm_slot;
|
|
uint16_t priority;
|
|
uint32_t waiting_timestamp;
|
|
bool is_waiting;
|
|
bool is_holding_lock;
|
|
|
|
ShmHelper *shm_helper;
|
|
std::thread heart_beat_thread;
|
|
std::atomic<bool> stop_heart_beat;
|
|
|
|
// request
|
|
uint64_t make_request(uint32_t type, uint64_t parameter);
|
|
int register_worker_shm();
|
|
|
|
// heart beat
|
|
void heart_beat_loop();
|
|
}; |