db/TaskletBlockingIO.cpp at main · carbonengine/db · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
// Copyright © 2024 CCP ehf.
#pragma once

#include "StdAfx.h"
#include "TaskletBlockingIO.h"
#include <BluePyCpp.h>
#include <Scheduler.h>
#include <socketmodule.h>

const void* COOKIE = "DB::IOWorkerContext";

static IOWorkerContext g_taskletBlockingRequestContext;

IOWorkerContext* GetTaskletBlockingRequestContext()
{
	return &g_taskletBlockingRequestContext;
}

IOWorkerContext::IOWorkerContext() :
	mLoop( nullptr )
{
}

bool IOWorkerContext::Init()
{
	auto* socketAPI = reinterpret_cast<PySocketModule_APIObject*>( PySocketModule_ImportModuleAndAPI() );
	if ( !socketAPI )
	{
		CCP_LOGERR( "Failed acquiring carbon-io socket module" );
		return false;
	}
	mLoop = socketAPI->get_uv_loop();
	if( !mLoop )
	{
		CCP_LOGERR( "Failed to get uv_loop from carbon-io socket module" );
		return false;
	}
	return true;
}

typedef struct {
	uv_work_t work;
	std::shared_ptr<IOWorker> request;
} uv_work_req_t;

void AfterWorkCB( uv_work_t* work, int status )
{
	auto workRequest = reinterpret_cast<uv_work_req_t*>(work);
	auto request = workRequest->request;
	if( status == UV_ECANCELED )
	{
		request->MarkCancelled();
	}
	else
	{
		request->Complete();
	}
	request.reset();
	delete workRequest;
}

void WorkCB( uv_work_t* work )
{
	auto request = reinterpret_cast<uv_work_req_t*>(work)->request;
	request->ThreadFunc();
}

void IOWorkerContext::Schedule( std::shared_ptr<IOWorker> request )
{
	auto* work = new uv_work_req_t;
	work->request = request;
	uv_queue_work( mLoop, reinterpret_cast<uv_work_t*>(work), WorkCB, AfterWorkCB );
}

IOWorker::IOWorker() :
	mState( IOWorker::PENDING )
{
	mChannel = SchedulerAPI()->PyChannel_New( nullptr );
	if( !mChannel )
	{
		PyErr_WriteUnraisable( PyUnicode_FromString( "IOWorker::Request Failed to create channel" ) );
		mState = IOWorker::FAILED;
	}
	SchedulerAPI()->PyChannel_SetPreference( mChannel, 1 );
}


IOWorker::~IOWorker()
{
	Py_XDECREF( mChannel );
	mChannel = nullptr;
}

bool IOWorker::ExecuteAndWait()
{
	g_taskletBlockingRequestContext.Schedule( shared_from_this() );
	auto sentinel = SchedulerAPI()->PyChannel_Receive( mChannel );
	if( !sentinel )
	{
		return false;
	}
	else
	{
		Py_DECREF( sentinel );
	}
	return true;
}

void IOWorker::MarkCancelled()
{
	mState = IOWorker::CANCELED;
	Ccp::PyGilEnsure gil;
	if( SchedulerAPI()->PyChannel_GetBalance( mChannel ) >= 0 )
	{
		return;
	}
	if( SchedulerAPI()->PyChannel_Send( mChannel, Py_None ) == -1 )
	{
		CCP_LOGWARN( "IOWorker::MarkCancelled Failed to send sentinel" );
	}
}

void IOWorker::Complete()
{
	mState = IOWorker::DONE;
	Ccp::PyGilEnsure gil;
	if( SchedulerAPI()->PyChannel_GetBalance( mChannel ) >= 0 )
	{
		return;
	}
	if( SchedulerAPI()->PyChannel_Send( mChannel, Py_None ) == -1 )
	{
		CCP_LOGWARN( "IOWorker::Complete Failed to send sentinel" );
	}
}