tudf.c 49.7 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/*
 * Copyright (c) 2019 TAOS Data, Inc. <jhtao@taosdata.com>
 *
 * This program is free software: you can use, redistribute, and/or modify
 * it under the terms of the GNU Affero General Public License, version 3
 * or later ("AGPL"), as published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 */
#include "uv.h"
#include "os.h"
S
slzhou 已提交
17
#include "fnLog.h"
18
#include "tudf.h"
19
#include "tudfInt.h"
S
shenglian zhou 已提交
20
#include "tarray.h"
S
slzhou 已提交
21
#include "tglobal.h"
S
slzhou 已提交
22
#include "tdatablock.h"
S
shenglian zhou 已提交
23 24
#include "querynodes.h"
#include "builtinsimpl.h"
S
slzhou 已提交
25
#include "functionMgt.h"
26

27
//TODO: add unit test
S
shenglian zhou 已提交
28
//TODO: include all global variable under context struct
S
slzhou 已提交
29

30 31 32 33 34 35 36
typedef struct SUdfdData {
  bool          startCalled;
  bool          needCleanUp;
  uv_loop_t     loop;
  uv_thread_t   thread;
  uv_barrier_t  barrier;
  uv_process_t  process;
37 38 39
#ifdef WINDOWS
  HANDLE        jobHandle;
#endif
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
  int           spawnErr;
  uv_pipe_t     ctrlPipe;
  uv_async_t    stopAsync;
  int32_t        stopCalled;

  int32_t         dnodeId;
} SUdfdData;

SUdfdData udfdGlobal = {0};

static int32_t udfSpawnUdfd(SUdfdData *pData);

void udfUdfdExit(uv_process_t *process, int64_t exitStatus, int termSignal) {
  fnInfo("udfd process exited with status %" PRId64 ", signal %d", exitStatus, termSignal);
  SUdfdData *pData = process->data;
  if (exitStatus == 0 && termSignal == 0 || atomic_load_32(&pData->stopCalled)) {
    fnInfo("udfd process exit due to SIGINT or dnode-mgmt called stop");
  } else {
    fnInfo("udfd process restart");
    udfSpawnUdfd(pData);
  }
}

static int32_t udfSpawnUdfd(SUdfdData* pData) {
  fnInfo("dnode start spawning udfd");
  uv_process_options_t options = {0};

  char path[PATH_MAX] = {0};
  if (tsProcPath == NULL) {
    path[0] = '.';
  } else {
    strncpy(path, tsProcPath, strlen(tsProcPath));
    taosDirName(path);
  }
#ifdef WINDOWS
  strcat(path, "udfd.exe");
#else
  strcat(path, "/udfd");
#endif
  char* argsUdfd[] = {path, "-c", configDir, NULL};
  options.args = argsUdfd;
  options.file = path;

  options.exit_cb = udfUdfdExit;

  uv_pipe_init(&pData->loop, &pData->ctrlPipe, 1);

  uv_stdio_container_t child_stdio[3];
  child_stdio[0].flags = UV_CREATE_PIPE | UV_READABLE_PIPE;
  child_stdio[0].data.stream = (uv_stream_t*) &pData->ctrlPipe;
  child_stdio[1].flags = UV_IGNORE;
  child_stdio[2].flags = UV_INHERIT_FD;
  child_stdio[2].data.fd = 2;
  options.stdio_count = 3;
  options.stdio = child_stdio;

  options.flags = UV_PROCESS_DETACHED;

  char dnodeIdEnvItem[32] = {0};
  char thrdPoolSizeEnvItem[32] = {0};
  snprintf(dnodeIdEnvItem, 32, "%s=%d", "DNODE_ID", pData->dnodeId);
  float numCpuCores = 4;
  taosGetCpuCores(&numCpuCores);
  snprintf(thrdPoolSizeEnvItem,32,  "%s=%d", "UV_THREADPOOL_SIZE", (int)numCpuCores*2);
  char* envUdfd[] = {dnodeIdEnvItem, thrdPoolSizeEnvItem, NULL};
  options.env = envUdfd;

  int err = uv_spawn(&pData->loop, &pData->process, &options);
  pData->process.data = (void*)pData;

110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127
#ifdef WINDOWS
  // End udfd.exe by Job.
  if (pData->jobHandle != NULL) CloseHandle(pData->jobHandle);
  pData->jobHandle = CreateJobObject(NULL, NULL);
  bool add_job_ok = AssignProcessToJobObject(pData->jobHandle, pData->process.process_handle);
  if (!add_job_ok) {
    fnError("Assign udfd to job failed.");
  } else {
    JOBOBJECT_EXTENDED_LIMIT_INFORMATION limit_info;
    memset(&limit_info, 0x0, sizeof(limit_info));
    limit_info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
    bool set_auto_kill_ok = SetInformationJobObject(pData->jobHandle, JobObjectExtendedLimitInformation, &limit_info, sizeof(limit_info));
    if (!set_auto_kill_ok) {
      fnError("Set job auto kill udfd failed.");
    }
  }
#endif

128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162
  if (err != 0) {
    fnError("can not spawn udfd. path: %s, error: %s", path, uv_strerror(err));
  }
  return err;
}

static void udfUdfdCloseWalkCb(uv_handle_t* handle, void* arg) {
  if (!uv_is_closing(handle)) {
    uv_close(handle, NULL);
  }
}

static void udfUdfdStopAsyncCb(uv_async_t *async) {
  SUdfdData *pData = async->data;
  uv_stop(&pData->loop);
}

static void udfWatchUdfd(void *args) {
  SUdfdData *pData = args;
  uv_loop_init(&pData->loop);
  uv_async_init(&pData->loop, &pData->stopAsync, udfUdfdStopAsyncCb);
  pData->stopAsync.data = pData;
  int32_t err = udfSpawnUdfd(pData);
  atomic_store_32(&pData->spawnErr, err);
  uv_barrier_wait(&pData->barrier);
  uv_run(&pData->loop, UV_RUN_DEFAULT);
  uv_loop_close(&pData->loop);

  uv_walk(&pData->loop, udfUdfdCloseWalkCb, NULL);
  uv_run(&pData->loop, UV_RUN_DEFAULT);
  uv_loop_close(&pData->loop);
  return;
}

int32_t udfStartUdfd(int32_t startDnodeId) {
S
slzhou 已提交
163 164
  if (!tsStartUdfd) {
    fnInfo("start udfd is disabled.")
S
slzhou 已提交
165
    return 0;
S
slzhou 已提交
166
  }
167 168
  SUdfdData *pData = &udfdGlobal;
  if (pData->startCalled) {
S
Shengliang Guan 已提交
169
    fnInfo("dnode start udfd already called");
170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186
    return 0;
  }
  pData->startCalled = true;
  char dnodeId[8] = {0};
  snprintf(dnodeId, sizeof(dnodeId), "%d", startDnodeId);
  uv_os_setenv("DNODE_ID", dnodeId);
  pData->dnodeId = startDnodeId;

  uv_barrier_init(&pData->barrier, 2);
  uv_thread_create(&pData->thread, udfWatchUdfd, pData);
  uv_barrier_wait(&pData->barrier);
  int32_t err = atomic_load_32(&pData->spawnErr);
  if (err != 0) {
    uv_barrier_destroy(&pData->barrier);
    uv_async_send(&pData->stopAsync);
    uv_thread_join(&pData->thread);
    pData->needCleanUp = false;
S
Shengliang Guan 已提交
187
    fnInfo("dnode udfd cleaned up after spawn err");
188 189 190 191 192 193 194 195
  } else {
    pData->needCleanUp = true;
  }
  return err;
}

int32_t udfStopUdfd() {
  SUdfdData *pData = &udfdGlobal;
S
Shengliang Guan 已提交
196
  fnInfo("dnode to stop udfd. need cleanup: %d, spawn err: %d",
197 198 199 200 201 202 203 204 205
        pData->needCleanUp, pData->spawnErr);
  if (!pData->needCleanUp || atomic_load_32(&pData->stopCalled)) {
    return 0;
  }
  atomic_store_32(&pData->stopCalled, 1);
  pData->needCleanUp = false;
  uv_barrier_destroy(&pData->barrier);
  uv_async_send(&pData->stopAsync);
  uv_thread_join(&pData->thread);
206 207 208
#ifdef WINDOWS
  if (pData->jobHandle != NULL) CloseHandle(pData->jobHandle);
#endif
S
Shengliang Guan 已提交
209
  fnInfo("dnode udfd cleaned up");
210 211 212 213
  return 0;
}

//==============================================================================================
S
shenglian zhou 已提交
214 215 216 217
/* Copyright (c) 2013, Ben Noordhuis <info@bnoordhuis.nl>
 * The QUEUE is copied from queue.h under libuv
 * */

S
shenglian zhou 已提交
218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305
typedef void *QUEUE[2];

/* Private macros. */
#define QUEUE_NEXT(q)       (*(QUEUE **) &((*(q))[0]))
#define QUEUE_PREV(q)       (*(QUEUE **) &((*(q))[1]))
#define QUEUE_PREV_NEXT(q)  (QUEUE_NEXT(QUEUE_PREV(q)))
#define QUEUE_NEXT_PREV(q)  (QUEUE_PREV(QUEUE_NEXT(q)))

/* Public macros. */
#define QUEUE_DATA(ptr, type, field)                                          \
  ((type *) ((char *) (ptr) - offsetof(type, field)))

/* Important note: mutating the list while QUEUE_FOREACH is
 * iterating over its elements results in undefined behavior.
 */
#define QUEUE_FOREACH(q, h)                                                   \
  for ((q) = QUEUE_NEXT(h); (q) != (h); (q) = QUEUE_NEXT(q))

#define QUEUE_EMPTY(q)                                                        \
  ((const QUEUE *) (q) == (const QUEUE *) QUEUE_NEXT(q))

#define QUEUE_HEAD(q)                                                         \
  (QUEUE_NEXT(q))

#define QUEUE_INIT(q)                                                         \
  do {                                                                        \
    QUEUE_NEXT(q) = (q);                                                      \
    QUEUE_PREV(q) = (q);                                                      \
  }                                                                           \
  while (0)

#define QUEUE_ADD(h, n)                                                       \
  do {                                                                        \
    QUEUE_PREV_NEXT(h) = QUEUE_NEXT(n);                                       \
    QUEUE_NEXT_PREV(n) = QUEUE_PREV(h);                                       \
    QUEUE_PREV(h) = QUEUE_PREV(n);                                            \
    QUEUE_PREV_NEXT(h) = (h);                                                 \
  }                                                                           \
  while (0)

#define QUEUE_SPLIT(h, q, n)                                                  \
  do {                                                                        \
    QUEUE_PREV(n) = QUEUE_PREV(h);                                            \
    QUEUE_PREV_NEXT(n) = (n);                                                 \
    QUEUE_NEXT(n) = (q);                                                      \
    QUEUE_PREV(h) = QUEUE_PREV(q);                                            \
    QUEUE_PREV_NEXT(h) = (h);                                                 \
    QUEUE_PREV(q) = (n);                                                      \
  }                                                                           \
  while (0)

#define QUEUE_MOVE(h, n)                                                      \
  do {                                                                        \
    if (QUEUE_EMPTY(h))                                                       \
      QUEUE_INIT(n);                                                          \
    else {                                                                    \
      QUEUE* q = QUEUE_HEAD(h);                                               \
      QUEUE_SPLIT(h, q, n);                                                   \
    }                                                                         \
  }                                                                           \
  while (0)

#define QUEUE_INSERT_HEAD(h, q)                                               \
  do {                                                                        \
    QUEUE_NEXT(q) = QUEUE_NEXT(h);                                            \
    QUEUE_PREV(q) = (h);                                                      \
    QUEUE_NEXT_PREV(q) = (q);                                                 \
    QUEUE_NEXT(h) = (q);                                                      \
  }                                                                           \
  while (0)

#define QUEUE_INSERT_TAIL(h, q)                                               \
  do {                                                                        \
    QUEUE_NEXT(q) = (h);                                                      \
    QUEUE_PREV(q) = QUEUE_PREV(h);                                            \
    QUEUE_PREV_NEXT(q) = (q);                                                 \
    QUEUE_PREV(h) = (q);                                                      \
  }                                                                           \
  while (0)

#define QUEUE_REMOVE(q)                                                       \
  do {                                                                        \
    QUEUE_PREV_NEXT(q) = QUEUE_NEXT(q);                                       \
    QUEUE_NEXT_PREV(q) = QUEUE_PREV(q);                                       \
  }                                                                           \
  while (0)


306 307 308 309 310 311
enum {
  UV_TASK_CONNECT = 0,
  UV_TASK_REQ_RSP = 1,
  UV_TASK_DISCONNECT = 2
};

312
int64_t gUdfTaskSeqNum = 0;
313
typedef struct SUdfcProxy {
314
  char udfdPipeName[PATH_MAX + UDF_LISTEN_PIPE_NAME_LEN + 2];
315
  uv_barrier_t initBarrier;
316

317 318 319
  uv_loop_t   uvLoop;
  uv_thread_t loopThread;
  uv_async_t  loopTaskAync;
320

321
  uv_async_t loopStopAsync;
322

323 324 325 326
  uv_mutex_t taskQueueMutex;
  int8_t     udfcState;
  QUEUE      taskQueue;
  QUEUE      uvProcTaskQueue;
327 328

  int8_t initialized;
329
} SUdfcProxy;
330

331
SUdfcProxy gUdfdProxy = {0};
332

S
slzhou 已提交
333
typedef struct SClientUdfUvSession {
334
  SUdfcProxy *udfc;
335
  int64_t severHandle;
S
slzhou 已提交
336
  uv_pipe_t *udfUvPipe;
S
shenglian zhou 已提交
337 338 339 340

  int8_t  outputType;
  int32_t outputLen;
  int32_t bufSize;
S
slzhou 已提交
341
} SClientUdfUvSession;
342 343

typedef struct SClientUvTaskNode {
344
  SUdfcProxy *udfc;
345 346 347 348 349 350 351 352 353 354 355
  int8_t type;
  int errCode;

  uv_pipe_t *pipe;

  int64_t seqNum;
  uv_buf_t reqBuf;

  uv_sem_t taskSem;
  uv_buf_t rspBuf;

S
shenglian zhou 已提交
356 357 358
  QUEUE recvTaskQueue;
  QUEUE procTaskQueue;
  QUEUE connTaskQueue;
359 360 361 362 363
} SClientUvTaskNode;

typedef struct SClientUdfTask {
  int8_t type;

S
slzhou 已提交
364
  SClientUdfUvSession *session;
365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393

  int32_t errCode;

  union {
    struct {
      SUdfSetupRequest req;
      SUdfSetupResponse rsp;
    } _setup;
    struct {
      SUdfCallRequest req;
      SUdfCallResponse rsp;
    } _call;
    struct {
      SUdfTeardownRequest req;
      SUdfTeardownResponse rsp;
    } _teardown;
  };

} SClientUdfTask;

typedef struct SClientConnBuf {
  char *buf;
  int32_t len;
  int32_t cap;
  int32_t total;
} SClientConnBuf;

typedef struct SClientUvConn {
  uv_pipe_t *pipe;
S
shenglian zhou 已提交
394
  QUEUE taskQueue;
395
  SClientConnBuf readBuf;
S
slzhou 已提交
396
  SClientUdfUvSession *session;
397 398
} SClientUvConn;

399 400
enum {
  UDFC_STATE_INITAL = 0, // initial state
401
  UDFC_STATE_STARTNG, // starting after udfcOpen
402
  UDFC_STATE_READY, // started and begin to receive quests
403
  UDFC_STATE_STOPPING, // stopping after udfcClose
404
};
405

406 407
int32_t getUdfdPipeName(char* pipeName, int32_t size) {
  char    dnodeId[8] = {0};
wafwerar's avatar
wafwerar 已提交
408
  size_t  dnodeIdSize = sizeof(dnodeId);
409 410
  int32_t err = uv_os_getenv(UDF_DNODE_ID_ENV_NAME, dnodeId, &dnodeIdSize);
  if (err != 0) {
411
    fnError("get dnode id from env. error: %s.", uv_err_name(err));
412 413
    dnodeId[0] = '1';
  }
414
#ifdef _WIN32
415
  snprintf(pipeName, size, "%s%s", UDF_LISTEN_PIPE_NAME_PREFIX, dnodeId);
416 417 418 419
#else
  snprintf(pipeName, size, "%s/%s%s", tsDataDir, UDF_LISTEN_PIPE_NAME_PREFIX, dnodeId);
#endif
  fnInfo("get dnode id from env. dnode id: %s. pipe path: %s", dnodeId, pipeName);
420 421 422
  return 0;
}

423 424 425
int32_t encodeUdfSetupRequest(void **buf, const SUdfSetupRequest *setup) {
  int32_t len = 0;
  len += taosEncodeBinary(buf, setup->udfName, TSDB_FUNC_NAME_LEN);
S
shenglian zhou 已提交
426 427
  return len;
}
428

429 430 431
void* decodeUdfSetupRequest(const void* buf, SUdfSetupRequest *request) {
  buf = taosDecodeBinaryTo(buf, request->udfName, TSDB_FUNC_NAME_LEN);
  return (void*)buf;
S
shenglian zhou 已提交
432
}
433

434 435
int32_t encodeUdfInterBuf(void **buf, const SUdfInterBuf* state) {
  int32_t len = 0;
436
  len += taosEncodeFixedI8(buf, state->numOfResult);
437 438
  len += taosEncodeFixedI32(buf, state->bufLen);
  len += taosEncodeBinary(buf, state->buf, state->bufLen);
S
shenglian zhou 已提交
439 440
  return len;
}
441

442
void* decodeUdfInterBuf(const void* buf, SUdfInterBuf* state) {
443
  buf = taosDecodeFixedI8(buf, &state->numOfResult);
444 445 446
  buf = taosDecodeFixedI32(buf, &state->bufLen);
  buf = taosDecodeBinary(buf, (void**)&state->buf, state->bufLen);
  return (void*)buf;
S
shenglian zhou 已提交
447 448
}

449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464
int32_t encodeUdfCallRequest(void **buf, const SUdfCallRequest *call) {
  int32_t len = 0;
  len += taosEncodeFixedI64(buf, call->udfHandle);
  len += taosEncodeFixedI8(buf, call->callType);
  if (call->callType == TSDB_UDF_CALL_SCALA_PROC) {
    len += tEncodeDataBlock(buf, &call->block);
  } else if (call->callType == TSDB_UDF_CALL_AGG_INIT) {
    len += taosEncodeFixedI8(buf, call->initFirst);
  } else if (call->callType == TSDB_UDF_CALL_AGG_PROC) {
    len += tEncodeDataBlock(buf, &call->block);
    len += encodeUdfInterBuf(buf, &call->interBuf);
  } else if (call->callType == TSDB_UDF_CALL_AGG_MERGE) {
    len += encodeUdfInterBuf(buf, &call->interBuf);
    len += encodeUdfInterBuf(buf, &call->interBuf2);
  } else if (call->callType == TSDB_UDF_CALL_AGG_FIN) {
    len += encodeUdfInterBuf(buf, &call->interBuf);
465
  }
466
  return len;
467 468
}

469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489
void* decodeUdfCallRequest(const void* buf, SUdfCallRequest* call) {
  buf = taosDecodeFixedI64(buf, &call->udfHandle);
  buf = taosDecodeFixedI8(buf, &call->callType);
  switch (call->callType) {
    case TSDB_UDF_CALL_SCALA_PROC:
      buf = tDecodeDataBlock(buf, &call->block);
      break;
    case TSDB_UDF_CALL_AGG_INIT:
      buf = taosDecodeFixedI8(buf, &call->initFirst);
      break;
    case TSDB_UDF_CALL_AGG_PROC:
      buf = tDecodeDataBlock(buf, &call->block);
      buf = decodeUdfInterBuf(buf, &call->interBuf);
      break;
    case TSDB_UDF_CALL_AGG_MERGE:
      buf = decodeUdfInterBuf(buf, &call->interBuf);
      buf = decodeUdfInterBuf(buf, &call->interBuf2);
      break;
    case TSDB_UDF_CALL_AGG_FIN:
      buf = decodeUdfInterBuf(buf, &call->interBuf);
      break;
490
  }
491
  return (void*)buf;
S
shenglian zhou 已提交
492 493
}

494 495 496 497
int32_t encodeUdfTeardownRequest(void **buf, const SUdfTeardownRequest *teardown) {
  int32_t len = 0;
  len += taosEncodeFixedI64(buf, teardown->udfHandle);
  return len;
S
shenglian zhou 已提交
498 499
}

500 501 502
void* decodeUdfTeardownRequest(const void* buf, SUdfTeardownRequest *teardown) {
  buf = taosDecodeFixedI64(buf, &teardown->udfHandle);
  return (void*)buf;
S
shenglian zhou 已提交
503 504
}

505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522
int32_t encodeUdfRequest(void** buf, const SUdfRequest* request) {
  int32_t len = 0;
  if (buf == NULL) {
    len += sizeof(request->msgLen);
  } else {
    *(int32_t*)(*buf) = request->msgLen;
    *buf = POINTER_SHIFT(*buf, sizeof(request->msgLen));
  }
  len += taosEncodeFixedI64(buf, request->seqNum);
  len += taosEncodeFixedI8(buf, request->type);
  if (request->type == UDF_TASK_SETUP) {
    len += encodeUdfSetupRequest(buf, &request->setup);
  } else if (request->type == UDF_TASK_CALL) {
    len += encodeUdfCallRequest(buf, &request->call);
  } else if (request->type == UDF_TASK_TEARDOWN) {
    len += encodeUdfTeardownRequest(buf, &request->teardown);
  }
  return len;
S
shenglian zhou 已提交
523 524
}

525 526
void* decodeUdfRequest(const void* buf, SUdfRequest* request) {
  request->msgLen = *(int32_t*)(buf);
S
slzhou 已提交
527
  buf = POINTER_SHIFT(buf, sizeof(request->msgLen));
S
shenglian zhou 已提交
528

529 530
  buf = taosDecodeFixedI64(buf, &request->seqNum);
  buf = taosDecodeFixedI8(buf, &request->type);
S
shenglian zhou 已提交
531 532

  if (request->type == UDF_TASK_SETUP) {
533
    buf = decodeUdfSetupRequest(buf, &request->setup);
S
shenglian zhou 已提交
534
  } else if (request->type == UDF_TASK_CALL) {
535 536 537
    buf = decodeUdfCallRequest(buf, &request->call);
  } else if (request->type == UDF_TASK_TEARDOWN) {
    buf = decodeUdfTeardownRequest(buf, &request->teardown);
S
shenglian zhou 已提交
538
  }
539
  return (void*)buf;
S
shenglian zhou 已提交
540
}
541

542 543 544
int32_t encodeUdfSetupResponse(void **buf, const SUdfSetupResponse *setupRsp) {
  int32_t len = 0;
  len += taosEncodeFixedI64(buf, setupRsp->udfHandle);
S
shenglian zhou 已提交
545 546 547
  len += taosEncodeFixedI8(buf, setupRsp->outputType);
  len += taosEncodeFixedI32(buf, setupRsp->outputLen);
  len += taosEncodeFixedI32(buf, setupRsp->bufSize);
548 549
  return len;
}
550

551 552
void* decodeUdfSetupResponse(const void* buf, SUdfSetupResponse* setupRsp) {
  buf = taosDecodeFixedI64(buf, &setupRsp->udfHandle);
S
shenglian zhou 已提交
553 554 555
  buf = taosDecodeFixedI8(buf, &setupRsp->outputType);
  buf = taosDecodeFixedI32(buf, &setupRsp->outputLen);
  buf = taosDecodeFixedI32(buf, &setupRsp->bufSize);
556
  return (void*)buf;
S
shenglian zhou 已提交
557
}
558

559 560 561 562 563 564
int32_t encodeUdfCallResponse(void **buf, const SUdfCallResponse *callRsp) {
  int32_t len = 0;
  len += taosEncodeFixedI8(buf, callRsp->callType);
  switch (callRsp->callType) {
    case TSDB_UDF_CALL_SCALA_PROC:
      len += tEncodeDataBlock(buf, &callRsp->resultData);
565
      break;
566
    case TSDB_UDF_CALL_AGG_INIT:
S
slzhou 已提交
567
      len += encodeUdfInterBuf(buf, &callRsp->resultBuf);
568 569
      break;
    case TSDB_UDF_CALL_AGG_PROC:
S
slzhou 已提交
570
      len += encodeUdfInterBuf(buf, &callRsp->resultBuf);
571 572
      break;
    case TSDB_UDF_CALL_AGG_MERGE:
S
slzhou 已提交
573
      len += encodeUdfInterBuf(buf, &callRsp->resultBuf);
574 575
      break;
    case TSDB_UDF_CALL_AGG_FIN:
S
slzhou 已提交
576
      len += encodeUdfInterBuf(buf, &callRsp->resultBuf);
577 578
      break;
  }
579
  return len;
S
shenglian zhou 已提交
580 581
}

582 583 584 585 586 587 588
void* decodeUdfCallResponse(const void* buf, SUdfCallResponse* callRsp) {
  buf = taosDecodeFixedI8(buf, &callRsp->callType);
  switch (callRsp->callType) {
    case TSDB_UDF_CALL_SCALA_PROC:
      buf = tDecodeDataBlock(buf, &callRsp->resultData);
      break;
    case TSDB_UDF_CALL_AGG_INIT:
S
slzhou 已提交
589
      buf = decodeUdfInterBuf(buf, &callRsp->resultBuf);
590 591
      break;
    case TSDB_UDF_CALL_AGG_PROC:
S
slzhou 已提交
592
      buf = decodeUdfInterBuf(buf, &callRsp->resultBuf);
593 594
      break;
    case TSDB_UDF_CALL_AGG_MERGE:
S
slzhou 已提交
595
      buf = decodeUdfInterBuf(buf, &callRsp->resultBuf);
596 597
      break;
    case TSDB_UDF_CALL_AGG_FIN:
S
slzhou 已提交
598
      buf = decodeUdfInterBuf(buf, &callRsp->resultBuf);
599
      break;
S
shenglian zhou 已提交
600
  }
601
  return (void*)buf;
S
shenglian zhou 已提交
602 603
}

604 605
int32_t encodeUdfTeardownResponse(void** buf, const SUdfTeardownResponse* teardownRsp) {
  return 0;
S
shenglian zhou 已提交
606 607
}

608 609
void* decodeUdfTeardownResponse(const void* buf, SUdfTeardownResponse* teardownResponse) {
  return (void*)buf;
S
shenglian zhou 已提交
610 611
}

612 613 614 615 616 617 618
int32_t encodeUdfResponse(void** buf, const SUdfResponse* rsp) {
  int32_t len = 0;
  if (buf == NULL) {
    len += sizeof(rsp->msgLen);
  } else {
    *(int32_t*)(*buf) = rsp->msgLen;
    *buf = POINTER_SHIFT(*buf, sizeof(rsp->msgLen));
S
shenglian zhou 已提交
619 620
  }

S
slzhou 已提交
621 622 623 624 625 626 627
  if (buf == NULL) {
    len += sizeof(rsp->seqNum);
  } else {
    *(int64_t*)(*buf) = rsp->seqNum;
    *buf = POINTER_SHIFT(*buf, sizeof(rsp->seqNum));
  }

628 629 630
  len += taosEncodeFixedI64(buf, rsp->seqNum);
  len += taosEncodeFixedI8(buf, rsp->type);
  len += taosEncodeFixedI32(buf, rsp->code);
S
shenglian zhou 已提交
631

632 633 634 635 636 637 638 639 640 641 642 643 644 645 646
  switch (rsp->type) {
    case UDF_TASK_SETUP:
      len += encodeUdfSetupResponse(buf, &rsp->setupRsp);
      break;
    case UDF_TASK_CALL:
      len += encodeUdfCallResponse(buf, &rsp->callRsp);
      break;
    case UDF_TASK_TEARDOWN:
      len += encodeUdfTeardownResponse(buf, &rsp->teardownRsp);
      break;
    default:
      //TODO: log error
      break;
  }
  return len;
S
shenglian zhou 已提交
647 648
}

649 650
void* decodeUdfResponse(const void* buf, SUdfResponse* rsp) {
  rsp->msgLen = *(int32_t*)(buf);
S
slzhou 已提交
651
  buf = POINTER_SHIFT(buf, sizeof(rsp->msgLen));
S
slzhou 已提交
652 653
  rsp->seqNum = *(int64_t*)(buf);
  buf = POINTER_SHIFT(buf, sizeof(rsp->seqNum));
654 655 656
  buf = taosDecodeFixedI64(buf, &rsp->seqNum);
  buf = taosDecodeFixedI8(buf, &rsp->type);
  buf = taosDecodeFixedI32(buf, &rsp->code);
S
shenglian zhou 已提交
657

658 659 660 661 662 663 664 665 666 667 668 669 670
  switch (rsp->type) {
    case UDF_TASK_SETUP:
      buf = decodeUdfSetupResponse(buf, &rsp->setupRsp);
      break;
    case UDF_TASK_CALL:
      buf = decodeUdfCallResponse(buf, &rsp->callRsp);
      break;
    case UDF_TASK_TEARDOWN:
      buf = decodeUdfTeardownResponse(buf, &rsp->teardownRsp);
      break;
    default:
      //TODO: log error
      break;
671
  }
672
  return (void*)buf;
673
}
674

S
shenglian zhou 已提交
675 676
void freeUdfColumnData(SUdfColumnData *data, SUdfColumnMeta *meta) {
  if (IS_VAR_DATA_TYPE(meta->type)) {
S
slzhou 已提交
677 678 679 680
    taosMemoryFree(data->varLenCol.varOffsets);
    data->varLenCol.varOffsets = NULL;
    taosMemoryFree(data->varLenCol.payload);
    data->varLenCol.payload = NULL;
S
shenglian zhou 已提交
681
  } else {
S
slzhou 已提交
682 683 684 685
    taosMemoryFree(data->fixLenCol.nullBitmap);
    data->fixLenCol.nullBitmap = NULL;
    taosMemoryFree(data->fixLenCol.data);
    data->fixLenCol.data = NULL;
S
shenglian zhou 已提交
686 687 688 689
  }
}

void freeUdfColumn(SUdfColumn* col) {
S
shenglian zhou 已提交
690
  freeUdfColumnData(&col->colData, &col->colMeta);
S
shenglian zhou 已提交
691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707
}

void freeUdfDataDataBlock(SUdfDataBlock *block) {
  for (int32_t i = 0; i < block->numOfCols; ++i) {
    freeUdfColumn(block->udfCols[i]);
    taosMemoryFree(block->udfCols[i]);
    block->udfCols[i] = NULL;
  }
  taosMemoryFree(block->udfCols);
  block->udfCols = NULL;
}

void freeUdfInterBuf(SUdfInterBuf *buf) {
  taosMemoryFree(buf->buf);
  buf->buf = NULL;
}

S
slzhou 已提交
708 709 710 711

int32_t convertDataBlockToUdfDataBlock(SSDataBlock *block, SUdfDataBlock *udfBlock) {
  udfBlock->numOfRows = block->info.rows;
  udfBlock->numOfCols = block->info.numOfCols;
S
slzhou 已提交
712
  udfBlock->udfCols = taosMemoryCalloc(udfBlock->numOfCols, sizeof(SUdfColumn*));
S
slzhou 已提交
713
  for (int32_t i = 0; i < udfBlock->numOfCols; ++i) {
S
slzhou 已提交
714
    udfBlock->udfCols[i] = taosMemoryCalloc(1, sizeof(SUdfColumn));
S
slzhou 已提交
715 716 717 718 719 720 721
    SColumnInfoData *col= (SColumnInfoData*)taosArrayGet(block->pDataBlock, i);
    SUdfColumn *udfCol = udfBlock->udfCols[i];
    udfCol->colMeta.type = col->info.type;
    udfCol->colMeta.bytes = col->info.bytes;
    udfCol->colMeta.scale = col->info.scale;
    udfCol->colMeta.precision = col->info.precision;
    udfCol->colData.numOfRows = udfBlock->numOfRows;
S
slzhou@taodata.com 已提交
722
    udfCol->hasNull = col->hasNull;
S
shenglian zhou 已提交
723
    if (IS_VAR_DATA_TYPE(udfCol->colMeta.type)) {
S
slzhou 已提交
724 725 726 727 728 729
      udfCol->colData.varLenCol.varOffsetsLen = sizeof(int32_t) * udfBlock->numOfRows;
      udfCol->colData.varLenCol.varOffsets = taosMemoryMalloc(udfCol->colData.varLenCol.varOffsetsLen);
      memcpy(udfCol->colData.varLenCol.varOffsets, col->varmeta.offset, udfCol->colData.varLenCol.varOffsetsLen);
      udfCol->colData.varLenCol.payloadLen = colDataGetLength(col, udfBlock->numOfRows);
      udfCol->colData.varLenCol.payload = taosMemoryMalloc(udfCol->colData.varLenCol.payloadLen);
      memcpy(udfCol->colData.varLenCol.payload, col->pData, udfCol->colData.varLenCol.payloadLen);
S
slzhou 已提交
730
    } else {
S
slzhou 已提交
731 732 733 734 735 736 737 738 739 740
      udfCol->colData.fixLenCol.nullBitmapLen = BitmapLen(udfCol->colData.numOfRows);
      int32_t bitmapLen = udfCol->colData.fixLenCol.nullBitmapLen;
      udfCol->colData.fixLenCol.nullBitmap = taosMemoryMalloc(udfCol->colData.fixLenCol.nullBitmapLen);
      char* bitmap = udfCol->colData.fixLenCol.nullBitmap;
      memcpy(bitmap, col->nullbitmap, bitmapLen);
      udfCol->colData.fixLenCol.dataLen = colDataGetLength(col, udfBlock->numOfRows);
      int32_t dataLen = udfCol->colData.fixLenCol.dataLen;
      udfCol->colData.fixLenCol.data = taosMemoryMalloc(udfCol->colData.fixLenCol.dataLen);
      char* data = udfCol->colData.fixLenCol.data;
      memcpy(data, col->pData, dataLen);
S
slzhou 已提交
741 742 743 744 745 746 747 748
    }
  }
  return 0;
}

int32_t convertUdfColumnToDataBlock(SUdfColumn *udfCol, SSDataBlock *block) {
  block->info.numOfCols = 1;
  block->info.rows = udfCol->colData.numOfRows;
S
shenglian zhou 已提交
749
  block->info.hasVarCol = IS_VAR_DATA_TYPE(udfCol->colMeta.type);
S
slzhou 已提交
750 751 752 753 754 755 756 757 758

  block->pDataBlock = taosArrayInit(1, sizeof(SColumnInfoData));
  taosArraySetSize(block->pDataBlock, 1);
  SColumnInfoData *col = taosArrayGet(block->pDataBlock, 0);
  SUdfColumnMeta *meta = &udfCol->colMeta;
  col->info.precision = meta->precision;
  col->info.bytes = meta->bytes;
  col->info.scale = meta->scale;
  col->info.type = meta->type;
S
slzhou@taodata.com 已提交
759
  col->hasNull = udfCol->hasNull;
S
slzhou 已提交
760 761 762
  SUdfColumnData *data = &udfCol->colData;

  if (!IS_VAR_DATA_TYPE(meta->type)) {
S
slzhou 已提交
763 764 765 766
    col->nullbitmap = taosMemoryMalloc(data->fixLenCol.nullBitmapLen);
    memcpy(col->nullbitmap, data->fixLenCol.nullBitmap, data->fixLenCol.nullBitmapLen);
    col->pData = taosMemoryMalloc(data->fixLenCol.dataLen);
    memcpy(col->pData, data->fixLenCol.data, data->fixLenCol.dataLen);
S
slzhou 已提交
767
  } else {
S
slzhou 已提交
768 769 770 771
    col->varmeta.offset = taosMemoryMalloc(data->varLenCol.varOffsetsLen);
    memcpy(col->varmeta.offset, data->varLenCol.varOffsets, data->varLenCol.varOffsetsLen);
    col->pData = taosMemoryMalloc(data->varLenCol.payloadLen);
    memcpy(col->pData, data->varLenCol.payload, data->varLenCol.payloadLen);
S
slzhou 已提交
772 773 774 775
  }
  return 0;
}

S
slzhou 已提交
776 777 778 779 780 781 782 783 784 785 786 787 788 789
int32_t convertScalarParamToDataBlock(SScalarParam *input, int32_t numOfCols, SSDataBlock *output) {
  output->info.rows = input->numOfRows;
  output->info.numOfCols = numOfCols;
  bool hasVarCol = false;
  for (int32_t i = 0; i < numOfCols; ++i) {
    if (IS_VAR_DATA_TYPE((input+i)->columnData->info.type)) {
      hasVarCol = true;
      break;
    }
  }
  output->info.hasVarCol = hasVarCol;

  //TODO: free the array output->pDataBlock
  output->pDataBlock = taosArrayInit(numOfCols, sizeof(SColumnInfoData));
790 791 792
  for (int32_t i = 0; i < numOfCols; ++i) {
    taosArrayPush(output->pDataBlock, (input + i)->columnData);
  }
S
slzhou 已提交
793 794 795 796 797 798 799 800 801 802 803 804 805
  return 0;
}

int32_t convertDataBlockToScalarParm(SSDataBlock *input, SScalarParam *output) {
  if (input->info.numOfCols != 1) {
    fnError("scalar function only support one column");
    return -1;
  }
  output->numOfRows = input->info.rows;
  //TODO: memory
  output->columnData = taosArrayGet(input->pDataBlock, 0);
  return 0;
}
S
slzhou 已提交
806

807 808
void onUdfcPipeClose(uv_handle_t *handle) {
  SClientUvConn *conn = handle->data;
S
shenglian zhou 已提交
809 810 811
  if (!QUEUE_EMPTY(&conn->taskQueue)) {
    QUEUE* h = QUEUE_HEAD(&conn->taskQueue);
    SClientUvTaskNode *task = QUEUE_DATA(h, SClientUvTaskNode, connTaskQueue);
812
    task->errCode = 0;
S
shenglian zhou 已提交
813
    QUEUE_REMOVE(&task->procTaskQueue);
S
slzhou 已提交
814
    uv_sem_post(&task->taskSem);
815
  }
S
slzhou 已提交
816
  conn->session->udfUvPipe = NULL;
wafwerar's avatar
wafwerar 已提交
817 818 819
  taosMemoryFree(conn->readBuf.buf);
  taosMemoryFree(conn);
  taosMemoryFree((uv_pipe_t *) handle);
820 821
}

S
slzhou 已提交
822 823
int32_t udfcGetUdfTaskResultFromUvTask(SClientUdfTask *task, SClientUvTaskNode *uvTask) {
  fnDebug("udfc get uv task result. task: %p, uvTask: %p", task, uvTask);
824 825
  if (uvTask->type == UV_TASK_REQ_RSP) {
    if (uvTask->rspBuf.base != NULL) {
S
shenglian zhou 已提交
826
      SUdfResponse rsp;
827 828
      void* buf = decodeUdfResponse(uvTask->rspBuf.base, &rsp);
      assert(uvTask->rspBuf.len == POINTER_DISTANCE(buf, uvTask->rspBuf.base));
S
shenglian zhou 已提交
829
      task->errCode = rsp.code;
830 831 832

      switch (task->type) {
        case UDF_TASK_SETUP: {
S
shenglian zhou 已提交
833
          //TODO: copy or not
S
shenglian zhou 已提交
834
          task->_setup.rsp = rsp.setupRsp;
835 836 837
          break;
        }
        case UDF_TASK_CALL: {
S
shenglian zhou 已提交
838
          task->_call.rsp = rsp.callRsp;
S
shenglian zhou 已提交
839
          //TODO: copy or not
840 841 842
          break;
        }
        case UDF_TASK_TEARDOWN: {
S
shenglian zhou 已提交
843
          task->_teardown.rsp = rsp.teardownRsp;
844 845 846 847 848 849 850 851 852
          //TODO: copy or not?
          break;
        }
        default: {
          break;
        }
      }

      // TODO: the call buffer is setup and freed by udf invocation
wafwerar's avatar
wafwerar 已提交
853
      taosMemoryFree(uvTask->rspBuf.base);
854 855
    } else {
      task->errCode = uvTask->errCode;
856
    }
857 858 859 860
  } else if (uvTask->type == UV_TASK_CONNECT) {
    task->errCode = uvTask->errCode;
  } else if (uvTask->type == UV_TASK_DISCONNECT) {
    task->errCode = uvTask->errCode;
861
  }
862 863 864 865 866 867 868 869 870
  return 0;
}

void udfcAllocateBuffer(uv_handle_t *handle, size_t suggestedSize, uv_buf_t *buf) {
  SClientUvConn *conn = handle->data;
  SClientConnBuf *connBuf = &conn->readBuf;

  int32_t msgHeadSize = sizeof(int32_t) + sizeof(int64_t);
  if (connBuf->cap == 0) {
wafwerar's avatar
wafwerar 已提交
871
    connBuf->buf = taosMemoryMalloc(msgHeadSize);
872 873 874 875 876 877 878 879
    if (connBuf->buf) {
      connBuf->len = 0;
      connBuf->cap = msgHeadSize;
      connBuf->total = -1;

      buf->base = connBuf->buf;
      buf->len = connBuf->cap;
    } else {
880
      fnError("udfc allocate buffer failure. size: %d", msgHeadSize);
881 882 883 884 885
      buf->base = NULL;
      buf->len = 0;
    }
  } else {
    connBuf->cap = connBuf->total > connBuf->cap ? connBuf->total : connBuf->cap;
wafwerar's avatar
wafwerar 已提交
886
    void *resultBuf = taosMemoryRealloc(connBuf->buf, connBuf->cap);
887 888 889 890 891
    if (resultBuf) {
      connBuf->buf = resultBuf;
      buf->base = connBuf->buf + connBuf->len;
      buf->len = connBuf->cap - connBuf->len;
    } else {
892
      fnError("udfc re-allocate buffer failure. size: %d", connBuf->cap);
893 894 895 896 897
      buf->base = NULL;
      buf->len = 0;
    }
  }

898
  fnTrace("conn buf cap - len - total : %d - %d - %d", connBuf->cap, connBuf->len, connBuf->total);
899 900 901

}

902 903 904 905 906
bool isUdfcUvMsgComplete(SClientConnBuf *connBuf) {
  if (connBuf->total == -1 && connBuf->len >= sizeof(int32_t)) {
    connBuf->total = *(int32_t *) (connBuf->buf);
  }
  if (connBuf->len == connBuf->cap && connBuf->total == connBuf->cap) {
907
    fnTrace("udfc complete message is received, now handle it");
908
    return true;
909
  }
910 911 912 913 914
  return false;
}

void udfcUvHandleRsp(SClientUvConn *conn) {
  SClientConnBuf *connBuf = &conn->readBuf;
915
  int64_t seqNum = *(int64_t *) (connBuf->buf + sizeof(int32_t)); // msglen then seqnum
916

S
shenglian zhou 已提交
917
  if (QUEUE_EMPTY(&conn->taskQueue)) {
918
    fnError("udfc no task waiting for response on connection");
919 920 921 922
    return;
  }
  bool found = false;
  SClientUvTaskNode *taskFound = NULL;
S
shenglian zhou 已提交
923 924 925 926
  QUEUE* h = QUEUE_NEXT(&conn->taskQueue);
  SClientUvTaskNode *task = QUEUE_DATA(h, SClientUvTaskNode, connTaskQueue);

  while (h != &conn->taskQueue) {
927 928 929 930 931
    if (task->seqNum == seqNum) {
      if (found == false) {
        found = true;
        taskFound = task;
      } else {
932
        fnError("udfc more than one task waiting for the same response");
933 934
        continue;
      }
935
    }
S
shenglian zhou 已提交
936 937
    h = QUEUE_NEXT(h);
    task = QUEUE_DATA(h, SClientUvTaskNode, connTaskQueue);
938 939
  }

940 941
  if (taskFound) {
    taskFound->rspBuf = uv_buf_init(connBuf->buf, connBuf->len);
S
shenglian zhou 已提交
942 943
    QUEUE_REMOVE(&taskFound->connTaskQueue);
    QUEUE_REMOVE(&taskFound->procTaskQueue);
S
slzhou 已提交
944
    uv_sem_post(&taskFound->taskSem);
945
  } else {
946
    fnError("no task is waiting for the response.");
947 948 949 950 951 952
  }
  connBuf->buf = NULL;
  connBuf->total = -1;
  connBuf->len = 0;
  connBuf->cap = 0;
}
953

954
void udfcUvHandleError(SClientUvConn *conn) {
S
shenglian zhou 已提交
955 956 957
  while (!QUEUE_EMPTY(&conn->taskQueue)) {
    QUEUE* h = QUEUE_HEAD(&conn->taskQueue);
    SClientUvTaskNode *task = QUEUE_DATA(h, SClientUvTaskNode, connTaskQueue);
958
    task->errCode = TSDB_CODE_UDF_PIPE_READ_ERR;
S
slzhou 已提交
959
    QUEUE_REMOVE(&task->connTaskQueue);
S
shenglian zhou 已提交
960
    QUEUE_REMOVE(&task->procTaskQueue);
S
slzhou 已提交
961
    uv_sem_post(&task->taskSem);
S
shenglian zhou 已提交
962 963
  }

S
slzhou 已提交
964
  uv_close((uv_handle_t *) conn->pipe, onUdfcPipeClose);
965 966 967
}

void onUdfcRead(uv_stream_t *client, ssize_t nread, const uv_buf_t *buf) {
968
  fnTrace("udfc client %p, client read from pipe. nread: %zd", client, nread);
969 970 971 972 973 974 975 976 977 978 979 980
  if (nread == 0) return;

  SClientUvConn *conn = client->data;
  SClientConnBuf *connBuf = &conn->readBuf;
  if (nread > 0) {
    connBuf->len += nread;
    if (isUdfcUvMsgComplete(connBuf)) {
      udfcUvHandleRsp(conn);
    }

  }
  if (nread < 0) {
S
slzhou 已提交
981
    fnError("udfc client pipe %p read error: %zd, %s.", client, nread, uv_strerror(nread));
982
    if (nread == UV_EOF) {
S
slzhou 已提交
983
      fnError("\tudfc client pipe %p closed", client);
984 985
    }
    udfcUvHandleError(conn);
986 987 988 989
  }

}

990 991
void onUdfClientWrite(uv_write_t *write, int status) {
  SClientUvTaskNode *uvTask = write->data;
992
  uv_pipe_t *pipe = uvTask->pipe;
993 994
  if (status == 0) {
    SClientUvConn *conn = pipe->data;
S
shenglian zhou 已提交
995
    QUEUE_INSERT_TAIL(&conn->taskQueue, &uvTask->connTaskQueue);
996
  } else {
997
    fnError("udfc client %p write error.", pipe);
998
  }
999
  fnTrace("udfc client %p write length:%zu", pipe, uvTask->reqBuf.len);
wafwerar's avatar
wafwerar 已提交
1000 1001
  taosMemoryFree(write);
  taosMemoryFree(uvTask->reqBuf.base);
1002
}
H
Haojun Liao 已提交
1003

1004 1005 1006 1007 1008
void onUdfClientConnect(uv_connect_t *connect, int status) {
  SClientUvTaskNode *uvTask = connect->data;
  uvTask->errCode = status;
  if (status != 0) {
    //TODO: LOG error
H
Haojun Liao 已提交
1009
  }
1010
  uv_read_start((uv_stream_t *) uvTask->pipe, udfcAllocateBuffer, onUdfcRead);
wafwerar's avatar
wafwerar 已提交
1011
  taosMemoryFree(connect);
1012
  uv_sem_post(&uvTask->taskSem);
S
shenglian zhou 已提交
1013
  QUEUE_REMOVE(&uvTask->procTaskQueue);
1014
}
H
Haojun Liao 已提交
1015

S
slzhou 已提交
1016
int32_t udfcCreateUvTask(SClientUdfTask *task, int8_t uvTaskType, SClientUvTaskNode **pUvTask) {
wafwerar's avatar
wafwerar 已提交
1017
  SClientUvTaskNode *uvTask = taosMemoryCalloc(1, sizeof(SClientUvTaskNode));
1018
  uvTask->type = uvTaskType;
1019
  uvTask->udfc = task->session->udfc;
1020 1021 1022

  if (uvTaskType == UV_TASK_CONNECT) {
  } else if (uvTaskType == UV_TASK_REQ_RSP) {
S
slzhou 已提交
1023
    uvTask->pipe = task->session->udfUvPipe;
1024 1025
    SUdfRequest request;
    request.type = task->type;
1026
    request.seqNum =   atomic_fetch_add_64(&gUdfTaskSeqNum, 1);
1027 1028

    if (task->type == UDF_TASK_SETUP) {
S
shenglian zhou 已提交
1029
      request.setup = task->_setup.req;
1030 1031
      request.type = UDF_TASK_SETUP;
    } else if (task->type == UDF_TASK_CALL) {
S
shenglian zhou 已提交
1032
      request.call = task->_call.req;
1033 1034
      request.type = UDF_TASK_CALL;
    } else if (task->type == UDF_TASK_TEARDOWN) {
S
shenglian zhou 已提交
1035
      request.teardown = task->_teardown.req;
1036 1037 1038 1039
      request.type = UDF_TASK_TEARDOWN;
    } else {
      //TODO log and return error
    }
1040 1041
    int32_t bufLen = encodeUdfRequest(NULL, &request);
    request.msgLen = bufLen;
S
slzhou 已提交
1042 1043
    void *bufBegin = taosMemoryMalloc(bufLen);
    void *buf = bufBegin;
1044
    encodeUdfRequest(&buf, &request);
S
slzhou 已提交
1045
    uvTask->reqBuf = uv_buf_init(bufBegin, bufLen);
1046 1047
    uvTask->seqNum = request.seqNum;
  } else if (uvTaskType == UV_TASK_DISCONNECT) {
S
slzhou 已提交
1048
    uvTask->pipe = task->session->udfUvPipe;
1049 1050
  }
  uv_sem_init(&uvTask->taskSem, 0);
H
Haojun Liao 已提交
1051

1052 1053 1054
  *pUvTask = uvTask;
  return 0;
}
H
Haojun Liao 已提交
1055

S
slzhou 已提交
1056
int32_t udfcQueueUvTask(SClientUvTaskNode *uvTask) {
1057
  fnTrace("queue uv task to event loop, task: %d, %p", uvTask->type, uvTask);
1058 1059 1060 1061 1062
  SUdfcProxy *udfc = uvTask->udfc;
  uv_mutex_lock(&udfc->taskQueueMutex);
  QUEUE_INSERT_TAIL(&udfc->taskQueue, &uvTask->recvTaskQueue);
  uv_mutex_unlock(&udfc->taskQueueMutex);
  uv_async_send(&udfc->loopTaskAync);
H
Haojun Liao 已提交
1063

1064
  uv_sem_wait(&uvTask->taskSem);
S
slzhou 已提交
1065
  fnInfo("udfc uv task finished. task: %d, %p", uvTask->type, uvTask);
1066
  uv_sem_destroy(&uvTask->taskSem);
H
Haojun Liao 已提交
1067

1068 1069
  return 0;
}
H
Haojun Liao 已提交
1070

S
slzhou 已提交
1071
int32_t udfcStartUvTask(SClientUvTaskNode *uvTask) {
1072
  fnTrace("event loop start uv task. task: %d, %p", uvTask->type, uvTask);
1073 1074
  switch (uvTask->type) {
    case UV_TASK_CONNECT: {
wafwerar's avatar
wafwerar 已提交
1075
      uv_pipe_t *pipe = taosMemoryMalloc(sizeof(uv_pipe_t));
1076
      uv_pipe_init(&uvTask->udfc->uvLoop, pipe, 0);
1077
      uvTask->pipe = pipe;
H
Haojun Liao 已提交
1078

S
slzhou 已提交
1079
      SClientUvConn *conn = taosMemoryCalloc(1, sizeof(SClientUvConn));
1080 1081 1082 1083 1084
      conn->pipe = pipe;
      conn->readBuf.len = 0;
      conn->readBuf.cap = 0;
      conn->readBuf.buf = 0;
      conn->readBuf.total = -1;
S
shenglian zhou 已提交
1085
      QUEUE_INIT(&conn->taskQueue);
H
Haojun Liao 已提交
1086

1087 1088
      pipe->data = conn;

wafwerar's avatar
wafwerar 已提交
1089
      uv_connect_t *connReq = taosMemoryMalloc(sizeof(uv_connect_t));
1090
      connReq->data = uvTask;
1091
      uv_pipe_connect(connReq, pipe, uvTask->udfc->udfdPipeName, onUdfClientConnect);
H
Haojun Liao 已提交
1092
      break;
1093 1094 1095
    }
    case UV_TASK_REQ_RSP: {
      uv_pipe_t *pipe = uvTask->pipe;
wafwerar's avatar
wafwerar 已提交
1096
      uv_write_t *write = taosMemoryMalloc(sizeof(uv_write_t));
1097 1098 1099 1100 1101 1102
      write->data = uvTask;
      uv_write(write, (uv_stream_t *) pipe, &uvTask->reqBuf, 1, onUdfClientWrite);
      break;
    }
    case UV_TASK_DISCONNECT: {
      SClientUvConn *conn = uvTask->pipe->data;
S
shenglian zhou 已提交
1103
      QUEUE_INSERT_TAIL(&conn->taskQueue, &uvTask->connTaskQueue);
1104 1105 1106 1107 1108 1109 1110
      uv_close((uv_handle_t *) uvTask->pipe, onUdfcPipeClose);
      break;
    }
    default: {
      break;
    }
  }
H
Haojun Liao 已提交
1111

1112 1113
  return 0;
}
H
Haojun Liao 已提交
1114

1115
void udfClientAsyncCb(uv_async_t *async) {
1116
  SUdfcProxy *udfc = async->data;
S
shenglian zhou 已提交
1117
  QUEUE wq;
1118

1119 1120 1121
  uv_mutex_lock(&udfc->taskQueueMutex);
  QUEUE_MOVE(&udfc->taskQueue, &wq);
  uv_mutex_unlock(&udfc->taskQueueMutex);
1122

S
shenglian zhou 已提交
1123 1124 1125 1126
  while (!QUEUE_EMPTY(&wq)) {
    QUEUE* h = QUEUE_HEAD(&wq);
    QUEUE_REMOVE(h);
    SClientUvTaskNode *task = QUEUE_DATA(h, SClientUvTaskNode, recvTaskQueue);
S
slzhou 已提交
1127
    udfcStartUvTask(task);
1128
    QUEUE_INSERT_TAIL(&udfc->uvProcTaskQueue, &task->procTaskQueue);
1129 1130 1131 1132
  }

}

1133
void cleanUpUvTasks(SUdfcProxy *udfc) {
S
slzhou 已提交
1134
  fnDebug("clean up uv tasks")
S
shenglian zhou 已提交
1135
  QUEUE wq;
1136

1137 1138 1139
  uv_mutex_lock(&udfc->taskQueueMutex);
  QUEUE_MOVE(&udfc->taskQueue, &wq);
  uv_mutex_unlock(&udfc->taskQueueMutex);
1140

S
shenglian zhou 已提交
1141 1142 1143 1144
  while (!QUEUE_EMPTY(&wq)) {
    QUEUE* h = QUEUE_HEAD(&wq);
    QUEUE_REMOVE(h);
    SClientUvTaskNode *task = QUEUE_DATA(h, SClientUvTaskNode, recvTaskQueue);
1145
    if (udfc->udfcState == UDFC_STATE_STOPPING) {
1146
      task->errCode = TSDB_CODE_UDF_STOPPING;
1147 1148 1149 1150
    }
    uv_sem_post(&task->taskSem);
  }

1151 1152
  while (!QUEUE_EMPTY(&udfc->uvProcTaskQueue)) {
    QUEUE* h = QUEUE_HEAD(&udfc->uvProcTaskQueue);
S
shenglian zhou 已提交
1153 1154
    QUEUE_REMOVE(h);
    SClientUvTaskNode *task = QUEUE_DATA(h, SClientUvTaskNode, procTaskQueue);
1155
    if (udfc->udfcState == UDFC_STATE_STOPPING) {
1156
      task->errCode = TSDB_CODE_UDF_STOPPING;
S
shenglian zhou 已提交
1157 1158 1159 1160
    }
    uv_sem_post(&task->taskSem);
  }
}
1161

S
shenglian zhou 已提交
1162
void udfStopAsyncCb(uv_async_t *async) {
1163
  SUdfcProxy *udfc = async->data;
1164
  cleanUpUvTasks(udfc);
1165 1166
  if (udfc->udfcState == UDFC_STATE_STOPPING) {
    uv_stop(&udfc->uvLoop);
S
shenglian zhou 已提交
1167
  }
1168
}
S
shenglian zhou 已提交
1169

S
shenglian zhou 已提交
1170
void constructUdfService(void *argsThread) {
1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181
  SUdfcProxy *udfc = (SUdfcProxy *)argsThread;
  uv_loop_init(&udfc->uvLoop);

  uv_async_init(&udfc->uvLoop, &udfc->loopTaskAync, udfClientAsyncCb);
  udfc->loopTaskAync.data = udfc;
  uv_async_init(&udfc->uvLoop, &udfc->loopStopAsync, udfStopAsyncCb);
  udfc->loopStopAsync.data = udfc;
  uv_mutex_init(&udfc->taskQueueMutex);
  QUEUE_INIT(&udfc->taskQueue);
  QUEUE_INIT(&udfc->uvProcTaskQueue);
  uv_barrier_wait(&udfc->initBarrier);
1182
  //TODO return value of uv_run
1183 1184
  uv_run(&udfc->uvLoop, UV_RUN_DEFAULT);
  uv_loop_close(&udfc->uvLoop);
1185 1186
}

1187 1188 1189 1190 1191
int32_t udfcOpen() {
  int8_t old = atomic_val_compare_exchange_8(&gUdfdProxy.initialized, 0, 1);
  if (old == 1) {
    return 0;
  }
1192
  SUdfcProxy *proxy = &gUdfdProxy;
1193
  getUdfdPipeName(proxy->udfdPipeName, sizeof(proxy->udfdPipeName));
1194 1195 1196 1197 1198 1199
  proxy->udfcState = UDFC_STATE_STARTNG;
  uv_barrier_init(&proxy->initBarrier, 2);
  uv_thread_create(&proxy->loopThread, constructUdfService, proxy);
  atomic_store_8(&proxy->udfcState, UDFC_STATE_READY);
  proxy->udfcState = UDFC_STATE_READY;
  uv_barrier_wait(&proxy->initBarrier);
1200
  fnInfo("udfc initialized")
1201 1202 1203
  return 0;
}

1204 1205 1206 1207 1208 1209
int32_t udfcClose() {
  int8_t old = atomic_val_compare_exchange_8(&gUdfdProxy.initialized, 1, 0);
  if (old == 0) {
    return 0;
  }

1210 1211 1212 1213 1214 1215 1216
  SUdfcProxy *udfc = &gUdfdProxy;
  udfc->udfcState = UDFC_STATE_STOPPING;
  uv_async_send(&udfc->loopStopAsync);
  uv_thread_join(&udfc->loopThread);
  uv_mutex_destroy(&udfc->taskQueueMutex);
  uv_barrier_destroy(&udfc->initBarrier);
  udfc->udfcState = UDFC_STATE_INITAL;
1217
  fnInfo("udfc cleaned up");
1218 1219 1220
  return 0;
}

S
slzhou 已提交
1221
int32_t udfcRunUdfUvTask(SClientUdfTask *task, int8_t uvTaskType) {
1222 1223
  SClientUvTaskNode *uvTask = NULL;

S
slzhou 已提交
1224 1225 1226
  udfcCreateUvTask(task, uvTaskType, &uvTask);
  udfcQueueUvTask(uvTask);
  udfcGetUdfTaskResultFromUvTask(task, uvTask);
1227
  if (uvTaskType == UV_TASK_CONNECT) {
S
slzhou 已提交
1228 1229 1230
    task->session->udfUvPipe = uvTask->pipe;
    SClientUvConn *conn = uvTask->pipe->data;
    conn->session = task->session;
S
slzhou 已提交
1231 1232
  }
  taosMemoryFree(uvTask);
1233 1234 1235 1236
  uvTask = NULL;
  return task->errCode;
}

1237 1238
int32_t setupUdf(char udfName[], UdfcFuncHandle *funcHandle) {
  fnInfo("udfc setup udf. udfName: %s", udfName);
1239
  if (gUdfdProxy.udfcState != UDFC_STATE_READY) {
1240
    return TSDB_CODE_UDF_INVALID_STATE;
1241
  }
S
slzhou 已提交
1242
  SClientUdfTask *task = taosMemoryCalloc(1,sizeof(SClientUdfTask));
1243
  task->errCode = 0;
S
slzhou 已提交
1244
  task->session = taosMemoryCalloc(1, sizeof(SClientUdfUvSession));
1245
  task->session->udfc = &gUdfdProxy;
1246 1247 1248
  task->type = UDF_TASK_SETUP;

  SUdfSetupRequest *req = &task->_setup.req;
S
shenglian zhou 已提交
1249
  memcpy(req->udfName, udfName, TSDB_FUNC_NAME_LEN);
1250

S
slzhou 已提交
1251
  int32_t errCode = udfcRunUdfUvTask(task, UV_TASK_CONNECT);
1252
  if (errCode != 0) {
1253
    fnError("failed to connect to pipe. udfName: %s, pipe: %s", udfName, (&gUdfdProxy)->udfdPipeName);
1254
    return TSDB_CODE_UDF_PIPE_CONNECT_ERR;
H
Haojun Liao 已提交
1255
  }
1256

S
slzhou 已提交
1257
  udfcRunUdfUvTask(task, UV_TASK_REQ_RSP);
1258 1259 1260

  SUdfSetupResponse *rsp = &task->_setup.rsp;
  task->session->severHandle = rsp->udfHandle;
S
shenglian zhou 已提交
1261 1262 1263
  task->session->outputType = rsp->outputType;
  task->session->outputLen = rsp->outputLen;
  task->session->bufSize = rsp->bufSize;
1264 1265 1266 1267 1268 1269
  if (task->errCode != 0) {
    fnError("failed to setup udf. err: %d", task->errCode)
  } else {
    fnInfo("sucessfully setup udf func handle. handle: %p", task->session);
    *funcHandle = task->session;
  }
1270
  int32_t err = task->errCode;
wafwerar's avatar
wafwerar 已提交
1271
  taosMemoryFree(task);
1272
  return err;
H
Haojun Liao 已提交
1273 1274
}

1275
int32_t callUdf(UdfcFuncHandle handle, int8_t callType, SSDataBlock *input, SUdfInterBuf *state, SUdfInterBuf *state2,
1276
                SSDataBlock* output, SUdfInterBuf *newState) {
1277
  fnTrace("udfc call udf. callType: %d, funcHandle: %p", callType, handle);
S
slzhou 已提交
1278 1279 1280
  SClientUdfUvSession *session = (SClientUdfUvSession *) handle;
  if (session->udfUvPipe == NULL) {
    fnError("No pipe to udfd");
1281
    return TSDB_CODE_UDF_PIPE_NO_PIPE;
S
slzhou 已提交
1282 1283
  }
  SClientUdfTask *task = taosMemoryCalloc(1, sizeof(SClientUdfTask));
1284
  task->errCode = 0;
S
slzhou 已提交
1285
  task->session = (SClientUdfUvSession *) handle;
1286 1287 1288
  task->type = UDF_TASK_CALL;

  SUdfCallRequest *req = &task->_call.req;
S
slzhou 已提交
1289
  req->udfHandle = task->session->severHandle;
S
slzhou 已提交
1290
  req->callType = callType;
S
slzhou 已提交
1291

S
shenglian zhou 已提交
1292
  switch (callType) {
1293 1294 1295 1296
    case TSDB_UDF_CALL_AGG_INIT: {
      req->initFirst = 1;
      break;
    }
S
shenglian zhou 已提交
1297 1298 1299 1300 1301
    case TSDB_UDF_CALL_AGG_PROC: {
      req->block = *input;
      req->interBuf = *state;
      break;
    }
1302 1303 1304 1305 1306 1307
    case TSDB_UDF_CALL_AGG_MERGE: {
      req->interBuf = *state;
      req->interBuf2 = *state2;
      break;
    }
    case TSDB_UDF_CALL_AGG_FIN: {
S
shenglian zhou 已提交
1308 1309 1310 1311 1312 1313 1314 1315 1316
      req->interBuf = *state;
      break;
    }
    case TSDB_UDF_CALL_SCALA_PROC: {
      req->block = *input;
      break;
    }
  }

S
slzhou 已提交
1317
  udfcRunUdfUvTask(task, UV_TASK_REQ_RSP);
1318

1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336 1337 1338 1339 1340 1341 1342 1343
  if (task->errCode != 0) {
    fnError("call udf failure. err: %d", task->errCode);
  } else {
    SUdfCallResponse *rsp = &task->_call.rsp;
    switch (callType) {
      case TSDB_UDF_CALL_AGG_INIT: {
        *newState = rsp->resultBuf;
        break;
      }
      case TSDB_UDF_CALL_AGG_PROC: {
        *newState = rsp->resultBuf;
        break;
      }
      case TSDB_UDF_CALL_AGG_MERGE: {
        *newState = rsp->resultBuf;
        break;
      }
      case TSDB_UDF_CALL_AGG_FIN: {
        *newState = rsp->resultBuf;
        break;
      }
      case TSDB_UDF_CALL_SCALA_PROC: {
        *output = rsp->resultData;
        break;
      }
S
shenglian zhou 已提交
1344
    }
S
slzhou 已提交
1345 1346
  };
  int err = task->errCode;
wafwerar's avatar
wafwerar 已提交
1347
  taosMemoryFree(task);
S
slzhou 已提交
1348
  return err;
1349 1350
}

1351
int32_t callUdfAggInit(UdfcFuncHandle handle, SUdfInterBuf *interBuf) {
S
slzhou 已提交
1352 1353 1354 1355 1356 1357 1358 1359 1360
  int8_t callType = TSDB_UDF_CALL_AGG_INIT;

  int32_t err = callUdf(handle, callType, NULL, NULL, NULL, NULL, interBuf);

  return err;
}

// input: block, state
// output: interbuf,
1361
int32_t callUdfAggProcess(UdfcFuncHandle handle, SSDataBlock *block, SUdfInterBuf *state, SUdfInterBuf *newState) {
S
slzhou 已提交
1362 1363 1364 1365 1366 1367 1368
  int8_t callType = TSDB_UDF_CALL_AGG_PROC;
  int32_t err = callUdf(handle, callType, block, state, NULL, NULL, newState);
  return err;
}

// input: interbuf1, interbuf2
// output: resultBuf
1369
int32_t callUdfAggMerge(UdfcFuncHandle handle, SUdfInterBuf *interBuf1, SUdfInterBuf *interBuf2, SUdfInterBuf *resultBuf) {
S
slzhou 已提交
1370 1371 1372 1373 1374 1375 1376
  int8_t callType = TSDB_UDF_CALL_AGG_MERGE;
  int32_t err = callUdf(handle, callType, NULL, interBuf1, interBuf2, NULL, resultBuf);
  return err;
}

// input: interBuf
// output: resultData
1377
int32_t callUdfAggFinalize(UdfcFuncHandle handle, SUdfInterBuf *interBuf, SUdfInterBuf *resultData) {
S
slzhou 已提交
1378
  int8_t callType = TSDB_UDF_CALL_AGG_FIN;
S
slzhou 已提交
1379 1380 1381 1382
  int32_t err = callUdf(handle, callType, NULL, interBuf, NULL, NULL, resultData);
  return err;
}

S
slzhou 已提交
1383
int32_t callUdfScalarFunc(UdfcFuncHandle handle, SScalarParam *input, int32_t numOfCols, SScalarParam* output) {
S
slzhou 已提交
1384
  int8_t callType = TSDB_UDF_CALL_SCALA_PROC;
S
slzhou 已提交
1385 1386 1387 1388
  SSDataBlock inputBlock = {0};
  convertScalarParamToDataBlock(input, numOfCols, &inputBlock);
  SSDataBlock resultBlock = {0};
  int32_t err = callUdf(handle, callType, &inputBlock, NULL, NULL, &resultBlock, NULL);
S
slzhou 已提交
1389 1390 1391
  if (err == 0) {
    convertDataBlockToScalarParm(&resultBlock, output);
  }
S
slzhou 已提交
1392 1393 1394
  return err;
}

1395
int32_t teardownUdf(UdfcFuncHandle handle) {
1396
  fnInfo("tear down udf. udf func handle: %p", handle);
1397

S
slzhou 已提交
1398 1399 1400
  SClientUdfUvSession *session = (SClientUdfUvSession *) handle;
  if (session->udfUvPipe == NULL) {
    fnError("pipe to udfd does not exist");
1401
    return TSDB_CODE_UDF_PIPE_NO_PIPE;
S
slzhou 已提交
1402 1403 1404
  }

  SClientUdfTask *task = taosMemoryCalloc(1, sizeof(SClientUdfTask));
1405
  task->errCode = 0;
S
slzhou 已提交
1406
  task->session = session;
1407 1408 1409 1410 1411
  task->type = UDF_TASK_TEARDOWN;

  SUdfTeardownRequest *req = &task->_teardown.req;
  req->udfHandle = task->session->severHandle;

S
slzhou 已提交
1412
  udfcRunUdfUvTask(task, UV_TASK_REQ_RSP);
1413 1414 1415 1416 1417

  SUdfTeardownResponse *rsp = &task->_teardown.rsp;

  int32_t err = task->errCode;

S
slzhou 已提交
1418
  udfcRunUdfUvTask(task, UV_TASK_DISCONNECT);
1419

wafwerar's avatar
wafwerar 已提交
1420 1421
  taosMemoryFree(task->session);
  taosMemoryFree(task);
1422 1423 1424

  return err;
}
S
shenglian zhou 已提交
1425

S
shenglian zhou 已提交
1426 1427
//memory layout |---SUdfAggRes----|-----final result-----|---inter result----|
typedef struct SUdfAggRes {
S
slzhou 已提交
1428
  SClientUdfUvSession *session;
S
shenglian zhou 已提交
1429 1430 1431 1432 1433 1434
  int8_t finalResNum;
  int8_t interResNum;
  char* finalResBuf;
  char* interResBuf;
} SUdfAggRes;

S
shenglian zhou 已提交
1435
bool udfAggGetEnv(struct SFunctionNode* pFunc, SFuncExecEnv* pEnv) {
S
slzhou 已提交
1436
  if (fmIsScalarFunc(pFunc->funcId)) {
S
shenglian zhou 已提交
1437 1438
    return false;
  }
S
slzhou 已提交
1439
  pEnv->calcMemSize = sizeof(SUdfAggRes) + pFunc->node.resType.bytes + pFunc->udfBufSize;
S
shenglian zhou 已提交
1440 1441 1442 1443 1444 1445 1446 1447
  return true;
}

bool udfAggInit(struct SqlFunctionCtx *pCtx, struct SResultRowEntryInfo* pResultCellInfo) {
  if (functionSetup(pCtx, pResultCellInfo) != true) {
    return false;
  }
  UdfcFuncHandle handle;
1448 1449 1450
  int32_t udfCode = 0;
  if ((udfCode = setupUdf((char*)pCtx->udfName, &handle)) != 0) {
    fnError("udfAggInit error. step setupUdf. udf code: %d", udfCode);
S
shenglian zhou 已提交
1451 1452
    return false;
  }
S
slzhou 已提交
1453
  SClientUdfUvSession *session = (SClientUdfUvSession *)handle;
S
shenglian zhou 已提交
1454 1455
  SUdfAggRes *udfRes = (SUdfAggRes*)GET_ROWCELL_INTERBUF(pResultCellInfo);
  int32_t envSize = sizeof(SUdfAggRes) + session->outputLen + session->bufSize;
S
shenglian zhou 已提交
1456
  memset(udfRes, 0, envSize);
S
shenglian zhou 已提交
1457

S
slzhou 已提交
1458 1459 1460
  udfRes->finalResBuf = (char*)udfRes + sizeof(SUdfAggRes);
  udfRes->interResBuf = (char*)udfRes + sizeof(SUdfAggRes) + session->outputLen;

S
slzhou 已提交
1461
  udfRes->session = (SClientUdfUvSession *)handle;
S
shenglian zhou 已提交
1462
  SUdfInterBuf buf = {0};
1463 1464
  if ((udfCode = callUdfAggInit(handle, &buf)) != 0) {
    fnError("udfAggInit error. step callUdfAggInit. udf code: %d", udfCode);
S
shenglian zhou 已提交
1465 1466
    return false;
  }
S
shenglian zhou 已提交
1467
  udfRes->interResNum = buf.numOfResult;
S
slzhou 已提交
1468
  memcpy(udfRes->interResBuf, buf.buf, buf.bufLen);
S
shenglian zhou 已提交
1469 1470 1471 1472 1473 1474 1475
  return true;
}

int32_t udfAggProcess(struct SqlFunctionCtx *pCtx) {
  SInputColumnInfoData* pInput = &pCtx->input;
  int32_t numOfCols = pInput->numOfInputCols;

S
slzhou 已提交
1476
  SUdfAggRes* udfRes = (SUdfAggRes *)GET_ROWCELL_INTERBUF(GET_RES_INFO(pCtx));
S
slzhou 已提交
1477
  SClientUdfUvSession *session = udfRes->session;
S
slzhou 已提交
1478 1479
  udfRes->finalResBuf = (char*)udfRes + sizeof(SUdfAggRes);
  udfRes->interResBuf = (char*)udfRes + sizeof(SUdfAggRes) + session->outputLen;
S
shenglian zhou 已提交
1480 1481 1482 1483 1484

  int32_t start = pInput->startRowIndex;
  int32_t numOfRows = pInput->numOfRows;


1485 1486
  SSDataBlock tempBlock = {0};
  tempBlock.info.numOfCols = numOfCols;
1487
  tempBlock.info.rows = pInput->totalRows;
1488
  tempBlock.info.uid = pInput->uid;
S
shenglian zhou 已提交
1489
  bool hasVarCol = false;
1490
  tempBlock.pDataBlock = taosArrayInit(numOfCols, sizeof(SColumnInfoData));
S
shenglian zhou 已提交
1491 1492 1493 1494 1495 1496

  for (int32_t i = 0; i < numOfCols; ++i) {
    SColumnInfoData *col = pInput->pData[i];
    if (IS_VAR_DATA_TYPE(col->info.type)) {
      hasVarCol = true;
    }
1497
    taosArrayPush(tempBlock.pDataBlock, col);
S
shenglian zhou 已提交
1498
  }
1499
  tempBlock.info.hasVarCol = hasVarCol;
S
shenglian zhou 已提交
1500

1501 1502
  SSDataBlock *inputBlock = blockDataExtractBlock(&tempBlock, start, numOfRows);

S
slzhou 已提交
1503
  SUdfInterBuf state = {.buf = udfRes->interResBuf,
1504
                        .bufLen = session->bufSize,
S
slzhou 已提交
1505
                        .numOfResult = udfRes->interResNum};
S
shenglian zhou 已提交
1506
  SUdfInterBuf newState = {0};
1507

1508 1509 1510 1511 1512 1513 1514 1515
  int32_t udfCode = callUdfAggProcess(session, inputBlock, &state, &newState);
  if (udfCode != 0) {
    fnError("udfAggProcess error. code: %d", udfCode);
    newState.numOfResult = 0;
  } else {
    udfRes->interResNum = newState.numOfResult;
    memcpy(udfRes->interResBuf, newState.buf, newState.bufLen);
  }
S
slzhou 已提交
1516 1517 1518
  if (newState.numOfResult == 1 || state.numOfResult == 1) {
    GET_RES_INFO(pCtx)->numOfRes = 1;
  }
S
shenglian zhou 已提交
1519

1520 1521
  blockDataDestroy(inputBlock);
  taosArrayDestroy(tempBlock.pDataBlock);
S
shenglian zhou 已提交
1522 1523

  taosMemoryFree(newState.buf);
1524
  return udfCode;
S
shenglian zhou 已提交
1525 1526 1527
}

int32_t udfAggFinalize(struct SqlFunctionCtx *pCtx, SSDataBlock* pBlock) {
S
slzhou 已提交
1528
  SUdfAggRes* udfRes = (SUdfAggRes *)GET_ROWCELL_INTERBUF(GET_RES_INFO(pCtx));
S
slzhou 已提交
1529
  SClientUdfUvSession *session = udfRes->session;
S
slzhou 已提交
1530 1531
  udfRes->finalResBuf = (char*)udfRes + sizeof(SUdfAggRes);
  udfRes->interResBuf = (char*)udfRes + sizeof(SUdfAggRes) + session->outputLen;
S
shenglian zhou 已提交
1532 1533


S
slzhou 已提交
1534
  SUdfInterBuf resultBuf = {0};
S
slzhou 已提交
1535
  SUdfInterBuf state = {.buf = udfRes->interResBuf,
1536
                        .bufLen = session->bufSize,
S
slzhou 已提交
1537
                        .numOfResult = udfRes->interResNum};
1538 1539 1540 1541 1542 1543 1544 1545 1546 1547
  int32_t udfCallCode= 0;
  udfCallCode= callUdfAggFinalize(session, &state, &resultBuf);
  if (udfCallCode!= 0) {
    fnError("udfAggFinalize error. callUdfAggFinalize step. udf code:%d", udfCallCode);
    GET_RES_INFO(pCtx)->numOfRes = 0;
  } else {
    memcpy(udfRes->finalResBuf, resultBuf.buf, session->outputLen);
    udfRes->finalResNum = resultBuf.numOfResult;
    GET_RES_INFO(pCtx)->numOfRes = udfRes->finalResNum;
  }
S
shenglian zhou 已提交
1548

1549 1550 1551
  int32_t code = teardownUdf(session);
  if (code != 0) {
    fnError("udfAggFinalize error. teardownUdf step. udf code: %d", code);
S
slzhou 已提交
1552
  }
1553

1554
  return functionFinalizeWithResultBuf(pCtx, pBlock, udfRes->finalResBuf);
1555

S
shenglian zhou 已提交
1556
}