clientTmq.c 102.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
/*
 * 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/>.
 */

L
Liu Jicong 已提交
16
#include "cJSON.h"
17 18 19
#include "clientInt.h"
#include "clientLog.h"
#include "parser.h"
H
Haojun Liao 已提交
20
#include "tdatablock.h"
21 22
#include "tdef.h"
#include "tglobal.h"
X
Xiaoyu Wang 已提交
23
#include "tqueue.h"
24
#include "tref.h"
L
Liu Jicong 已提交
25 26
#include "ttimer.h"

X
Xiaoyu Wang 已提交
27 28
#define EMPTY_BLOCK_POLL_IDLE_DURATION 10
#define DEFAULT_AUTO_COMMIT_INTERVAL   5000
29

30 31
#define OFFSET_IS_RESET_OFFSET(_of)  ((_of) < 0)

32 33
typedef void (*__tmq_askep_fn_t)(tmq_t* pTmq, int32_t code, SDataBuf* pBuf, void* pParam);

X
Xiaoyu Wang 已提交
34
struct SMqMgmt {
35 36 37
  int8_t  inited;
  tmr_h   timer;
  int32_t rsetId;
38
};
L
Liu Jicong 已提交
39

X
Xiaoyu Wang 已提交
40 41
static TdThreadOnce   tmqInit = PTHREAD_ONCE_INIT;  // initialize only once
volatile int32_t      tmqInitRes = 0;               // initialize rsp code
42
static struct SMqMgmt tmqMgmt = {0};
43

L
Liu Jicong 已提交
44 45 46 47 48 49
typedef struct {
  int8_t  tmqRspType;
  int32_t epoch;
} SMqRspWrapper;

typedef struct {
L
Liu Jicong 已提交
50 51 52
  int8_t      tmqRspType;
  int32_t     epoch;
  SMqAskEpRsp msg;
L
Liu Jicong 已提交
53 54
} SMqAskEpRspWrapper;

L
Liu Jicong 已提交
55
struct tmq_list_t {
L
Liu Jicong 已提交
56
  SArray container;
L
Liu Jicong 已提交
57
};
L
Liu Jicong 已提交
58

L
Liu Jicong 已提交
59
struct tmq_conf_t {
60 61 62 63 64 65 66 67
  char           clientId[256];
  char           groupId[TSDB_CGROUP_LEN];
  int8_t         autoCommit;
  int8_t         resetOffset;
  int8_t         withTbName;
  int8_t         snapEnable;
  int32_t        snapBatchSize;
  bool           hbBgEnable;
68 69 70 71 72
  uint16_t       port;
  int32_t        autoCommitInterval;
  char*          ip;
  char*          user;
  char*          pass;
73
  tmq_commit_cb* commitCb;
L
Liu Jicong 已提交
74
  void*          commitCbUserParam;
L
Liu Jicong 已提交
75 76 77
};

struct tmq_t {
78 79 80 81 82 83 84
  int64_t        refId;
  char           groupId[TSDB_CGROUP_LEN];
  char           clientId[256];
  int8_t         withTbName;
  int8_t         useSnapshot;
  int8_t         autoCommit;
  int32_t        autoCommitInterval;
85
  int8_t         resetOffsetCfg;
86 87
  uint64_t       consumerId;
  bool           hbBgEnable;
L
Liu Jicong 已提交
88 89
  tmq_commit_cb* commitCb;
  void*          commitCbUserParam;
L
Liu Jicong 已提交
90 91

  // status
wmmhello's avatar
wmmhello 已提交
92
  SRWLatch        lock;
L
Liu Jicong 已提交
93 94
  int8_t  status;
  int32_t epoch;
L
Liu Jicong 已提交
95 96
#if 0
  int8_t  epStatus;
L
Liu Jicong 已提交
97
  int32_t epSkipCnt;
L
Liu Jicong 已提交
98
#endif
99
  // poll info
X
Xiaoyu Wang 已提交
100 101
  int64_t pollCnt;
  int64_t totalRows;
wmmhello's avatar
wmmhello 已提交
102
//  bool    needReportOffsetRows;
L
Liu Jicong 已提交
103

L
Liu Jicong 已提交
104
  // timer
X
Xiaoyu Wang 已提交
105 106 107 108 109 110 111 112 113 114
  tmr_h       hbLiveTimer;
  tmr_h       epTimer;
  tmr_h       reportTimer;
  tmr_h       commitTimer;
  STscObj*    pTscObj;       // connection
  SArray*     clientTopics;  // SArray<SMqClientTopic>
  STaosQueue* mqueue;        // queue of rsp
  STaosQall*  qall;
  STaosQueue* delayedTask;  // delayed task queue for heartbeat and auto commit
  tsem_t      rspSem;
L
Liu Jicong 已提交
115 116
};

117 118
typedef struct SAskEpInfo {
  int32_t code;
H
Haojun Liao 已提交
119
  tsem_t  sem;
120 121
} SAskEpInfo;

X
Xiaoyu Wang 已提交
122 123 124 125 126 127 128 129
enum {
  TMQ_VG_STATUS__IDLE = 0,
  TMQ_VG_STATUS__WAIT,
};

enum {
  TMQ_CONSUMER_STATUS__INIT = 0,
  TMQ_CONSUMER_STATUS__READY,
130
  TMQ_CONSUMER_STATUS__NO_TOPIC,
L
Liu Jicong 已提交
131
  TMQ_CONSUMER_STATUS__RECOVER,
L
Liu Jicong 已提交
132 133
};

L
Liu Jicong 已提交
134
enum {
135
  TMQ_DELAYED_TASK__ASK_EP = 1,
L
Liu Jicong 已提交
136 137 138 139
  TMQ_DELAYED_TASK__REPORT,
  TMQ_DELAYED_TASK__COMMIT,
};

H
Haojun Liao 已提交
140
typedef struct SVgOffsetInfo {
L
Liu Jicong 已提交
141
  STqOffsetVal committedOffset;
142 143
  STqOffsetVal endOffset;        // the last version in TAOS_RES + 1
  STqOffsetVal beginOffset;      // the first version in TAOS_RES
H
Haojun Liao 已提交
144 145 146 147 148 149 150 151 152 153
  int64_t      walVerBegin;
  int64_t      walVerEnd;
} SVgOffsetInfo;

typedef struct {
  int64_t       pollCnt;
  int64_t       numOfRows;
  SVgOffsetInfo offsetInfo;
  int32_t       vgId;
  int32_t       vgStatus;
H
Haojun Liao 已提交
154
  int32_t       vgSkipCnt;              // here used to mark the slow vgroups
155
//  bool          receivedInfoFromVnode;  // has already received info from vnode
H
Haojun Liao 已提交
156 157
  int64_t       emptyBlockReceiveTs;    // once empty block is received, idle for ignoreCnt then start to poll data
  bool          seekUpdated;            // offset is updated by seek operator, therefore, not update by vnode rsp.
H
Haojun Liao 已提交
158
  SEpSet        epSet;
159 160
} SMqClientVg;

L
Liu Jicong 已提交
161
typedef struct {
162 163 164
  char           topicName[TSDB_TOPIC_FNAME_LEN];
  char           db[TSDB_DB_FNAME_LEN];
  SArray*        vgs;  // SArray<SMqClientVg>
L
Liu Jicong 已提交
165
  SSchemaWrapper schema;
166 167
} SMqClientTopic;

L
Liu Jicong 已提交
168 169
typedef struct {
  int8_t          tmqRspType;
170
  int32_t         epoch;  // epoch can be used to guard the vgHandle
171
  int32_t         vgId;
wmmhello's avatar
wmmhello 已提交
172
  char            topicName[TSDB_TOPIC_FNAME_LEN];
L
Liu Jicong 已提交
173 174
  SMqClientVg*    vgHandle;
  SMqClientTopic* topicHandle;
H
Haojun Liao 已提交
175
  uint64_t        reqId;
176
  SEpSet*         pEpset;
L
Liu Jicong 已提交
177
  union {
L
Liu Jicong 已提交
178 179
    SMqDataRsp dataRsp;
    SMqMetaRsp metaRsp;
L
Liu Jicong 已提交
180
    STaosxRsp  taosxRsp;
L
Liu Jicong 已提交
181
  };
L
Liu Jicong 已提交
182 183
} SMqPollRspWrapper;

L
Liu Jicong 已提交
184
typedef struct {
wmmhello's avatar
wmmhello 已提交
185 186
//  int64_t refId;
//  int32_t epoch;
L
Liu Jicong 已提交
187 188
  tsem_t  rspSem;
  int32_t rspErr;
L
Liu Jicong 已提交
189
} SMqSubscribeCbParam;
L
Liu Jicong 已提交
190

L
Liu Jicong 已提交
191
typedef struct {
192 193 194 195
  int64_t          refId;
  int32_t          epoch;
  void*            pParam;
  __tmq_askep_fn_t pUserFn;
196 197
} SMqAskEpCbParam;

L
Liu Jicong 已提交
198
typedef struct {
199 200
  int64_t         refId;
  int32_t         epoch;
wmmhello's avatar
wmmhello 已提交
201 202 203
  char            topicName[TSDB_TOPIC_FNAME_LEN];
//  SMqClientVg*    pVg;
//  SMqClientTopic* pTopic;
L
Liu Jicong 已提交
204
  int32_t         vgId;
X
Xiaoyu Wang 已提交
205
  uint64_t        requestId;  // request id for debug purpose
X
Xiaoyu Wang 已提交
206
} SMqPollCbParam;
207

208 209 210 211 212 213 214 215 216 217
typedef struct SMqVgCommon {
  tsem_t        rsp;
  int32_t       numOfRsp;
  SArray*       pList;
  TdThreadMutex mutex;
  int64_t       consumerId;
  char*         pTopicName;
  int32_t       code;
} SMqVgCommon;

218 219 220 221 222
typedef struct SMqSeekParam {
  tsem_t        sem;
  int32_t       code;
} SMqSeekParam;

223 224 225 226 227 228
typedef struct SMqCommittedParam {
  tsem_t        sem;
  int32_t       code;
  SMqVgOffset   vgOffset;
} SMqCommittedParam;

229 230 231 232 233 234 235
typedef struct SMqVgWalInfoParam {
  int32_t      vgId;
  int32_t      epoch;
  int32_t      totalReq;
  SMqVgCommon* pCommon;
} SMqVgWalInfoParam;

236
typedef struct {
237 238
  int64_t        refId;
  int32_t        epoch;
L
Liu Jicong 已提交
239 240
  int32_t        waitingRspNum;
  int32_t        totalRspNum;
241
  int32_t        code;
242
  tmq_commit_cb* callbackFn;
L
Liu Jicong 已提交
243 244
  /*SArray*        successfulOffsets;*/
  /*SArray*        failedOffsets;*/
X
Xiaoyu Wang 已提交
245
  void* userParam;
246 247 248 249
} SMqCommitCbParamSet;

typedef struct {
  SMqCommitCbParamSet* params;
250
//  SMqVgOffset*         pOffset;
H
Haojun Liao 已提交
251 252 253
  char                 topicName[TSDB_TOPIC_FNAME_LEN];
  int32_t              vgId;
  tmq_t*               pTmq;
254
} SMqCommitCbParam;
255

256 257 258 259 260
typedef struct SSyncCommitInfo {
  tsem_t  sem;
  int32_t code;
} SSyncCommitInfo;

261
static int32_t doAskEp(tmq_t* tmq);
262 263
static int32_t makeTopicVgroupKey(char* dst, const char* topicName, int32_t vg);
static int32_t tmqCommitDone(SMqCommitCbParamSet* pParamSet);
264
static int32_t doSendCommitMsg(tmq_t* tmq, int32_t vgId, SEpSet* epSet, STqOffsetVal* offset, const char* pTopicName, SMqCommitCbParamSet* pParamSet);
X
Xiaoyu Wang 已提交
265 266 267
static void    commitRspCountDown(SMqCommitCbParamSet* pParamSet, int64_t consumerId, const char* pTopic, int32_t vgId);
static void    asyncAskEp(tmq_t* pTmq, __tmq_askep_fn_t askEpFn, void* param);
static void    addToQueueCallbackFn(tmq_t* pTmq, int32_t code, SDataBuf* pDataBuf, void* param);
268

269
tmq_conf_t* tmq_conf_new() {
wafwerar's avatar
wafwerar 已提交
270
  tmq_conf_t* conf = taosMemoryCalloc(1, sizeof(tmq_conf_t));
271 272 273 274 275
  if (conf == NULL) {
    terrno = TSDB_CODE_OUT_OF_MEMORY;
    return conf;
  }

276
  conf->withTbName = false;
L
Liu Jicong 已提交
277
  conf->autoCommit = true;
278
  conf->autoCommitInterval = DEFAULT_AUTO_COMMIT_INTERVAL;
279
  conf->resetOffset = TMQ_OFFSET__RESET_EARLIEST;
280
  conf->hbBgEnable = true;
281

282 283 284
  return conf;
}

L
Liu Jicong 已提交
285
void tmq_conf_destroy(tmq_conf_t* conf) {
L
Liu Jicong 已提交
286
  if (conf) {
287 288 289 290 291 292 293 294 295
    if (conf->ip) {
      taosMemoryFree(conf->ip);
    }
    if (conf->user) {
      taosMemoryFree(conf->user);
    }
    if (conf->pass) {
      taosMemoryFree(conf->pass);
    }
L
Liu Jicong 已提交
296 297
    taosMemoryFree(conf);
  }
L
Liu Jicong 已提交
298 299 300
}

tmq_conf_res_t tmq_conf_set(tmq_conf_t* conf, const char* key, const char* value) {
301 302 303
  if (conf == NULL || key == NULL || value == NULL){
    return TMQ_CONF_INVALID;
  }
304
  if (strcasecmp(key, "group.id") == 0) {
L
Liu Jicong 已提交
305
    tstrncpy(conf->groupId, value, TSDB_CGROUP_LEN);
L
Liu Jicong 已提交
306
    return TMQ_CONF_OK;
307
  }
L
Liu Jicong 已提交
308

309
  if (strcasecmp(key, "client.id") == 0) {
L
Liu Jicong 已提交
310
    tstrncpy(conf->clientId, value, 256);
L
Liu Jicong 已提交
311 312
    return TMQ_CONF_OK;
  }
L
Liu Jicong 已提交
313

314 315
  if (strcasecmp(key, "enable.auto.commit") == 0) {
    if (strcasecmp(value, "true") == 0) {
L
Liu Jicong 已提交
316
      conf->autoCommit = true;
L
Liu Jicong 已提交
317
      return TMQ_CONF_OK;
318
    } else if (strcasecmp(value, "false") == 0) {
L
Liu Jicong 已提交
319
      conf->autoCommit = false;
L
Liu Jicong 已提交
320 321 322 323
      return TMQ_CONF_OK;
    } else {
      return TMQ_CONF_INVALID;
    }
324
  }
L
Liu Jicong 已提交
325

326
  if (strcasecmp(key, "auto.commit.interval.ms") == 0) {
327
    conf->autoCommitInterval = taosStr2int64(value);
L
Liu Jicong 已提交
328 329 330
    return TMQ_CONF_OK;
  }

331 332 333
  if (strcasecmp(key, "auto.offset.reset") == 0) {
    if (strcasecmp(value, "none") == 0) {
      conf->resetOffset = TMQ_OFFSET__RESET_NONE;
L
Liu Jicong 已提交
334
      return TMQ_CONF_OK;
335
    } else if (strcasecmp(value, "earliest") == 0) {
336
      conf->resetOffset = TMQ_OFFSET__RESET_EARLIEST;
L
Liu Jicong 已提交
337
      return TMQ_CONF_OK;
338 339
    } else if (strcasecmp(value, "latest") == 0) {
      conf->resetOffset = TMQ_OFFSET__RESET_LATEST;
L
Liu Jicong 已提交
340 341 342 343 344
      return TMQ_CONF_OK;
    } else {
      return TMQ_CONF_INVALID;
    }
  }
L
Liu Jicong 已提交
345

346 347
  if (strcasecmp(key, "msg.with.table.name") == 0) {
    if (strcasecmp(value, "true") == 0) {
348
      conf->withTbName = true;
L
Liu Jicong 已提交
349
      return TMQ_CONF_OK;
350
    } else if (strcasecmp(value, "false") == 0) {
351
      conf->withTbName = false;
L
Liu Jicong 已提交
352
      return TMQ_CONF_OK;
353 354 355 356 357
    } else {
      return TMQ_CONF_INVALID;
    }
  }

358 359
  if (strcasecmp(key, "experimental.snapshot.enable") == 0) {
    if (strcasecmp(value, "true") == 0) {
L
Liu Jicong 已提交
360
      conf->snapEnable = true;
361
      return TMQ_CONF_OK;
362
    } else if (strcasecmp(value, "false") == 0) {
L
Liu Jicong 已提交
363
      conf->snapEnable = false;
364 365 366 367 368 369
      return TMQ_CONF_OK;
    } else {
      return TMQ_CONF_INVALID;
    }
  }

370
  if (strcasecmp(key, "experimental.snapshot.batch.size") == 0) {
371
    conf->snapBatchSize = taosStr2int64(value);
L
Liu Jicong 已提交
372 373 374
    return TMQ_CONF_OK;
  }

375
//  if (strcasecmp(key, "enable.heartbeat.background") == 0) {
X
Xiaoyu Wang 已提交
376 377 378 379 380 381 382
    //    if (strcasecmp(value, "true") == 0) {
    //      conf->hbBgEnable = true;
    //      return TMQ_CONF_OK;
    //    } else if (strcasecmp(value, "false") == 0) {
    //      conf->hbBgEnable = false;
    //      return TMQ_CONF_OK;
    //    } else {
383 384
//    tscError("the default value of enable.heartbeat.background is true, can not be seted");
//    return TMQ_CONF_INVALID;
X
Xiaoyu Wang 已提交
385
    //    }
386
//  }
L
Liu Jicong 已提交
387

388
  if (strcasecmp(key, "td.connect.ip") == 0) {
389
    conf->ip = taosStrdup(value);
L
Liu Jicong 已提交
390 391
    return TMQ_CONF_OK;
  }
392

393
  if (strcasecmp(key, "td.connect.user") == 0) {
394
    conf->user = taosStrdup(value);
L
Liu Jicong 已提交
395 396
    return TMQ_CONF_OK;
  }
397

398
  if (strcasecmp(key, "td.connect.pass") == 0) {
399
    conf->pass = taosStrdup(value);
L
Liu Jicong 已提交
400 401
    return TMQ_CONF_OK;
  }
402

403
  if (strcasecmp(key, "td.connect.port") == 0) {
404
    conf->port = taosStr2int64(value);
L
Liu Jicong 已提交
405 406
    return TMQ_CONF_OK;
  }
407

408
  if (strcasecmp(key, "td.connect.db") == 0) {
L
Liu Jicong 已提交
409 410 411
    return TMQ_CONF_OK;
  }

L
Liu Jicong 已提交
412
  return TMQ_CONF_UNKNOWN;
413 414
}

X
Xiaoyu Wang 已提交
415
tmq_list_t* tmq_list_new() { return (tmq_list_t*)taosArrayInit(0, sizeof(void*)); }
416

L
Liu Jicong 已提交
417
int32_t tmq_list_append(tmq_list_t* list, const char* src) {
418
  if(list == NULL) return -1;
L
Liu Jicong 已提交
419
  SArray* container = &list->container;
420
  if (src == NULL || src[0] == 0) return -1;
421
  char* topic = taosStrdup(src);
L
fix  
Liu Jicong 已提交
422
  if (taosArrayPush(container, &topic) == NULL) return -1;
423 424 425
  return 0;
}

L
Liu Jicong 已提交
426
void tmq_list_destroy(tmq_list_t* list) {
427
  if(list == NULL) return;
L
Liu Jicong 已提交
428
  SArray* container = &list->container;
L
Liu Jicong 已提交
429
  taosArrayDestroyP(container, taosMemoryFree);
L
Liu Jicong 已提交
430 431
}

L
Liu Jicong 已提交
432
int32_t tmq_list_get_size(const tmq_list_t* list) {
433
  if(list == NULL) return -1;
L
Liu Jicong 已提交
434 435 436 437 438
  const SArray* container = &list->container;
  return taosArrayGetSize(container);
}

char** tmq_list_to_c_array(const tmq_list_t* list) {
439
  if(list == NULL) return NULL;
L
Liu Jicong 已提交
440 441 442 443
  const SArray* container = &list->container;
  return container->pData;
}

H
Haojun Liao 已提交
444
static int32_t tmqCommitCb(void* param, SDataBuf* pBuf, int32_t code) {
445
  SMqCommitCbParam*    pParam = (SMqCommitCbParam*)param;
446
  SMqCommitCbParamSet* pParamSet = (SMqCommitCbParamSet*)pParam->params;
H
Haojun Liao 已提交
447

448
//  taosMemoryFree(pParam->pOffset);
L
Liu Jicong 已提交
449
  taosMemoryFree(pBuf->pData);
dengyihao's avatar
dengyihao 已提交
450
  taosMemoryFree(pBuf->pEpSet);
L
Liu Jicong 已提交
451

452
  commitRspCountDown(pParamSet, pParam->pTmq->consumerId, pParam->topicName, pParam->vgId);
453 454 455
  return 0;
}

456
static int32_t doSendCommitMsg(tmq_t* tmq, int32_t vgId, SEpSet* epSet, STqOffsetVal* offset, const char* pTopicName, SMqCommitCbParamSet* pParamSet) {
457
  SMqVgOffset pOffset = {0};
458

459 460
  pOffset.consumerId = tmq->consumerId;
  pOffset.offset.val = *offset;
461

L
Liu Jicong 已提交
462
  int32_t groupLen = strlen(tmq->groupId);
463 464 465
  memcpy(pOffset.offset.subKey, tmq->groupId, groupLen);
  pOffset.offset.subKey[groupLen] = TMQ_SEPARATOR;
  strcpy(pOffset.offset.subKey + groupLen + 1, pTopicName);
L
Liu Jicong 已提交
466

467 468
  int32_t len = 0;
  int32_t code = 0;
469
  tEncodeSize(tEncodeMqVgOffset, &pOffset, len, code);
L
Liu Jicong 已提交
470
  if (code < 0) {
471
    return TSDB_CODE_INVALID_PARA;
L
Liu Jicong 已提交
472
  }
473

L
Liu Jicong 已提交
474
  void* buf = taosMemoryCalloc(1, sizeof(SMsgHead) + len);
L
Liu Jicong 已提交
475
  if (buf == NULL) {
476
    return TSDB_CODE_OUT_OF_MEMORY;
L
Liu Jicong 已提交
477
  }
478

479
  ((SMsgHead*)buf)->vgId = htonl(vgId);
L
Liu Jicong 已提交
480

L
Liu Jicong 已提交
481 482 483 484
  void* abuf = POINTER_SHIFT(buf, sizeof(SMsgHead));

  SEncoder encoder;
  tEncoderInit(&encoder, abuf, len);
485
  tEncodeMqVgOffset(&encoder, &pOffset);
L
Liu Jicong 已提交
486
  tEncoderClear(&encoder);
L
Liu Jicong 已提交
487 488

  // build param
489
  SMqCommitCbParam* pParam = taosMemoryCalloc(1, sizeof(SMqCommitCbParam));
L
Liu Jicong 已提交
490 491
  if (pParam == NULL) {
    taosMemoryFree(buf);
492
    return TSDB_CODE_OUT_OF_MEMORY;
L
Liu Jicong 已提交
493
  }
494

L
Liu Jicong 已提交
495
  pParam->params = pParamSet;
496
//  pParam->pOffset = pOffset;
497
  pParam->vgId = vgId;
H
Haojun Liao 已提交
498 499
  pParam->pTmq = tmq;

H
Haojun Liao 已提交
500
  tstrncpy(pParam->topicName, pTopicName, tListLen(pParam->topicName));
L
Liu Jicong 已提交
501 502 503 504

  // build send info
  SMsgSendInfo* pMsgSendInfo = taosMemoryCalloc(1, sizeof(SMsgSendInfo));
  if (pMsgSendInfo == NULL) {
L
Liu Jicong 已提交
505 506
    taosMemoryFree(buf);
    taosMemoryFree(pParam);
507
    return TSDB_CODE_OUT_OF_MEMORY;
L
Liu Jicong 已提交
508
  }
509

510
  pMsgSendInfo->msgInfo = (SDataBuf) { .pData = buf, .len = sizeof(SMsgHead) + len, .handle = NULL };
L
Liu Jicong 已提交
511 512 513 514

  pMsgSendInfo->requestId = generateRequestId();
  pMsgSendInfo->requestObjRefId = 0;
  pMsgSendInfo->param = pParam;
L
Liu Jicong 已提交
515
  pMsgSendInfo->paramFreeFp = taosMemoryFree;
516
  pMsgSendInfo->fp = tmqCommitCb;
517
  pMsgSendInfo->msgType = TDMT_VND_TMQ_COMMIT_OFFSET;
L
Liu Jicong 已提交
518

L
Liu Jicong 已提交
519 520 521
  atomic_add_fetch_32(&pParamSet->waitingRspNum, 1);
  atomic_add_fetch_32(&pParamSet->totalRspNum, 1);

522
  SEp* pEp = GET_ACTIVE_EP(epSet);
523

H
Haojun Liao 已提交
524

L
Liu Jicong 已提交
525
  int64_t transporterId = 0;
526
  return asyncSendMsgToServer(tmq->pTscObj->pAppInfo->pTransporter, epSet, &transporterId, pMsgSendInfo);
L
Liu Jicong 已提交
527 528
}

H
Haojun Liao 已提交
529 530 531 532 533 534 535 536 537 538 539
static SMqClientTopic* getTopicByName(tmq_t* tmq, const char* pTopicName) {
  int32_t numOfTopics = taosArrayGetSize(tmq->clientTopics);
  for (int32_t i = 0; i < numOfTopics; ++i) {
    SMqClientTopic* pTopic = taosArrayGet(tmq->clientTopics, i);
    if (strcmp(pTopic->topicName, pTopicName) != 0) {
      continue;
    }

    return pTopic;
  }

H
Haojun Liao 已提交
540
  tscError("consumer:0x%" PRIx64 ", total:%d, failed to find topic:%s", tmq->consumerId, numOfTopics, pTopicName);
H
Haojun Liao 已提交
541 542 543
  return NULL;
}

544
static SMqCommitCbParamSet* prepareCommitCbParamSet(tmq_t* tmq, tmq_commit_cb* pCommitFp, void* userParam, int32_t rspNum){
L
Liu Jicong 已提交
545 546
  SMqCommitCbParamSet* pParamSet = taosMemoryCalloc(1, sizeof(SMqCommitCbParamSet));
  if (pParamSet == NULL) {
547
    return NULL;
L
Liu Jicong 已提交
548
  }
H
Haojun Liao 已提交
549

550 551
  pParamSet->refId = tmq->refId;
  pParamSet->epoch = tmq->epoch;
552
  pParamSet->callbackFn = pCommitFp;
L
Liu Jicong 已提交
553
  pParamSet->userParam = userParam;
554
  pParamSet->waitingRspNum = rspNum;
L
Liu Jicong 已提交
555

556 557
  return pParamSet;
}
558

559 560 561


static int32_t getClientVg(tmq_t* tmq, char* pTopicName, int32_t vgId, SMqClientVg** pVg){
H
Haojun Liao 已提交
562 563
  SMqClientTopic* pTopic = getTopicByName(tmq, pTopicName);
  if (pTopic == NULL) {
564 565
    tscError("consumer:0x%" PRIx64 " invalid topic name:%s", tmq->consumerId, pTopicName);
    return TSDB_CODE_TMQ_INVALID_TOPIC;
566
  }
L
Liu Jicong 已提交
567

568 569 570 571 572
  int32_t  numOfVgs = taosArrayGetSize(pTopic->vgs);
  for (int32_t i = 0; i < numOfVgs; ++i) {
    SMqClientVg* pClientVg = taosArrayGet(pTopic->vgs, i);
    if (pClientVg->vgId == vgId) {
      *pVg = pClientVg;
573
      break;
L
Liu Jicong 已提交
574
    }
L
Liu Jicong 已提交
575
  }
L
Liu Jicong 已提交
576

577
  return *pVg == NULL ? TSDB_CODE_TMQ_INVALID_VGID : TSDB_CODE_SUCCESS;
578 579 580 581 582 583
}

static int32_t asyncCommitOffset(tmq_t* tmq, char* pTopicName, int32_t vgId, STqOffsetVal* offsetVal, tmq_commit_cb* pCommitFp, void* userParam) {
  int32_t code = 0;
  tscInfo("consumer:0x%" PRIx64 " do manual commit offset for %s, vgId:%d", tmq->consumerId, pTopicName, vgId);
  taosRLockLatch(&tmq->lock);
584 585 586
  SMqClientVg* pVg = NULL;
  code = getClientVg(tmq, pTopicName, vgId, &pVg);
  if(code != 0){
587 588 589 590 591 592 593 594
    goto end;
  }
  if (offsetVal->type > 0 && !tOffsetEqual(offsetVal, &pVg->offsetInfo.committedOffset)) {
    char offsetBuf[TSDB_OFFSET_LEN] = {0};
    tFormatOffset(offsetBuf, tListLen(offsetBuf), offsetVal);

    char commitBuf[TSDB_OFFSET_LEN] = {0};
    tFormatOffset(commitBuf, tListLen(commitBuf), &pVg->offsetInfo.committedOffset);
L
Liu Jicong 已提交
595

596 597 598 599 600
    SMqCommitCbParamSet* pParamSet = prepareCommitCbParamSet(tmq, pCommitFp, userParam, 0);
    if (pParamSet == NULL) {
      code = TSDB_CODE_OUT_OF_MEMORY;
      goto end;
    }
601
    code = doSendCommitMsg(tmq, pVg->vgId, &pVg->epSet, offsetVal, pTopicName, pParamSet);
602
    if (code != TSDB_CODE_SUCCESS) {
603 604
      tscError("consumer:0x%" PRIx64 " topic:%s on vgId:%d end commit msg failed, send offset:%s committed:%s, code:%s",
               tmq->consumerId, pTopicName, pVg->vgId, offsetBuf, commitBuf, tstrerror(terrno));
605
      taosMemoryFree(pParamSet);
606
      goto end;
607
    }
608 609 610 611

    tscInfo("consumer:0x%" PRIx64 " topic:%s on vgId:%d send commit msg success, send offset:%s committed:%s",
            tmq->consumerId, pTopicName, pVg->vgId, offsetBuf, commitBuf);
    pVg->offsetInfo.committedOffset = *offsetVal;
L
Liu Jicong 已提交
612
  }
613 614

end:
wmmhello's avatar
wmmhello 已提交
615
  taosRUnLockLatch(&tmq->lock);
616
  return code;
L
Liu Jicong 已提交
617 618
}

619 620 621 622 623 624 625 626 627
static void asyncCommitFromResult(tmq_t* tmq, const TAOS_RES* pRes, tmq_commit_cb* pCommitFp, void* userParam){
  char*   pTopicName = NULL;
  int32_t vgId = 0;
  STqOffsetVal offsetVal = {0};
  int32_t code = 0;

  if (pRes == NULL || tmq == NULL) {
    code = TSDB_CODE_INVALID_PARA;
    goto end;
628
  }
629

630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652
  if (TD_RES_TMQ(pRes)) {
    SMqRspObj* pRspObj = (SMqRspObj*)pRes;
    pTopicName = pRspObj->topic;
    vgId = pRspObj->vgId;
    offsetVal = pRspObj->rsp.rspOffset;
  } else if (TD_RES_TMQ_META(pRes)) {
    SMqMetaRspObj* pMetaRspObj = (SMqMetaRspObj*)pRes;
    pTopicName = pMetaRspObj->topic;
    vgId = pMetaRspObj->vgId;
    offsetVal = pMetaRspObj->metaRsp.rspOffset;
  } else if (TD_RES_TMQ_METADATA(pRes)) {
    SMqTaosxRspObj* pRspObj = (SMqTaosxRspObj*)pRes;
    pTopicName = pRspObj->topic;
    vgId = pRspObj->vgId;
    offsetVal = pRspObj->rsp.rspOffset;
  } else {
    code = TSDB_CODE_TMQ_INVALID_MSG;
    goto end;
  }

  code = asyncCommitOffset(tmq, pTopicName, vgId, &offsetVal, pCommitFp, userParam);

end:
653
  if(code != TSDB_CODE_SUCCESS && pCommitFp != NULL){
654 655 656
    pCommitFp(tmq, code, userParam);
  }
}
657

658 659
static void asyncCommitAllOffsets(tmq_t* tmq, tmq_commit_cb* pCommitFp, void* userParam) {
  int32_t code = 0;
660
  // init as 1 to prevent concurrency issue
661 662 663 664 665
  SMqCommitCbParamSet* pParamSet = prepareCommitCbParamSet(tmq, pCommitFp, userParam, 1);
  if (pParamSet == NULL) {
    code = TSDB_CODE_OUT_OF_MEMORY;
    goto end;
  }
666

wmmhello's avatar
wmmhello 已提交
667
  taosRLockLatch(&tmq->lock);
668
  int32_t numOfTopics = taosArrayGetSize(tmq->clientTopics);
669
  tscInfo("consumer:0x%" PRIx64 " start to commit offset for %d topics", tmq->consumerId, numOfTopics);
670 671

  for (int32_t i = 0; i < numOfTopics; i++) {
672
    SMqClientTopic* pTopic = taosArrayGet(tmq->clientTopics, i);
673
    int32_t         numOfVgroups = taosArrayGetSize(pTopic->vgs);
L
Liu Jicong 已提交
674

675
    tscInfo("consumer:0x%" PRIx64 " commit offset for topics:%s, numOfVgs:%d", tmq->consumerId, pTopic->topicName, numOfVgroups);
676
    for (int32_t j = 0; j < numOfVgroups; j++) {
677 678
      SMqClientVg* pVg = taosArrayGet(pTopic->vgs, j);

679 680 681 682 683 684 685 686
      if (pVg->offsetInfo.endOffset.type > 0 && !tOffsetEqual(&pVg->offsetInfo.endOffset, &pVg->offsetInfo.committedOffset)) {
        char offsetBuf[TSDB_OFFSET_LEN] = {0};
        tFormatOffset(offsetBuf, tListLen(offsetBuf), &pVg->offsetInfo.endOffset);

        char commitBuf[TSDB_OFFSET_LEN] = {0};
        tFormatOffset(commitBuf, tListLen(commitBuf), &pVg->offsetInfo.committedOffset);

        code = doSendCommitMsg(tmq, pVg->vgId, &pVg->epSet, &pVg->offsetInfo.endOffset, pTopic->topicName, pParamSet);
687
        if (code != TSDB_CODE_SUCCESS) {
688 689
          tscError("consumer:0x%" PRIx64 " topic:%s on vgId:%d end commit msg failed, send offset:%s committed:%s, code:%s ordinal:%d/%d",
                   tmq->consumerId, pTopic->topicName, pVg->vgId, offsetBuf, commitBuf, tstrerror(terrno), j + 1, numOfVgroups);
L
Liu Jicong 已提交
690 691
          continue;
        }
H
Haojun Liao 已提交
692

693 694 695
        tscInfo("consumer:0x%" PRIx64 " topic:%s on vgId:%d send commit msg success, send offset:%s committed:%s, ordinal:%d/%d",
                tmq->consumerId, pTopic->topicName, pVg->vgId, offsetBuf, commitBuf, j + 1, numOfVgroups);
        pVg->offsetInfo.committedOffset = pVg->offsetInfo.endOffset;
696
      } else {
697 698
        tscInfo("consumer:0x%" PRIx64 " topic:%s vgId:%d, no commit, current:%" PRId64 ", ordinal:%d/%d",
                 tmq->consumerId, pTopic->topicName, pVg->vgId, pVg->offsetInfo.endOffset.version, j + 1, numOfVgroups);
699 700 701
      }
    }
  }
wmmhello's avatar
wmmhello 已提交
702
  taosRUnLockLatch(&tmq->lock);
703

704
  tscInfo("consumer:0x%" PRIx64 " total commit:%d for %d topics", tmq->consumerId, pParamSet->waitingRspNum - 1, numOfTopics);
H
Haojun Liao 已提交
705

706 707 708 709
  // request is sent
  if (pParamSet->totalRspNum != 0) {
    // count down since waiting rsp num init as 1
    commitRspCountDown(pParamSet, tmq->consumerId, "", 0);
710
    return;
L
Liu Jicong 已提交
711 712
  }

713 714
end:
  taosMemoryFree(pParamSet);
715
  if(pCommitFp != NULL) {
716 717
    pCommitFp(tmq, code, userParam);
  }
718
  return;
719 720
}

721 722
static void generateTimedTask(int64_t refId, int32_t type) {
  tmq_t* tmq = taosAcquireRef(tmqMgmt.rsetId, refId);
wmmhello's avatar
wmmhello 已提交
723 724 725 726 727 728 729 730 731
  if(tmq == NULL) return;

  int8_t* pTaskType = taosAllocateQitem(sizeof(int8_t), DEF_QITEM, 0);
  if(pTaskType == NULL) return;

  *pTaskType = type;
  taosWriteQitem(tmq->delayedTask, pTaskType);
  tsem_post(&tmq->rspSem);
  taosReleaseRef(tmqMgmt.rsetId, refId);
732 733 734 735 736
}

void tmqAssignAskEpTask(void* param, void* tmrId) {
  int64_t refId = *(int64_t*)param;
  generateTimedTask(refId, TMQ_DELAYED_TASK__ASK_EP);
737
  taosMemoryFree(param);
L
Liu Jicong 已提交
738 739 740
}

void tmqAssignDelayedCommitTask(void* param, void* tmrId) {
741
  int64_t refId = *(int64_t*)param;
742
  generateTimedTask(refId, TMQ_DELAYED_TASK__COMMIT);
743
  taosMemoryFree(param);
L
Liu Jicong 已提交
744 745
}

wmmhello's avatar
wmmhello 已提交
746 747 748 749 750 751 752 753 754 755 756 757 758
//void tmqAssignDelayedReportTask(void* param, void* tmrId) {
//  int64_t refId = *(int64_t*)param;
//  tmq_t*  tmq = taosAcquireRef(tmqMgmt.rsetId, refId);
//  if (tmq != NULL) {
//    int8_t* pTaskType = taosAllocateQitem(sizeof(int8_t), DEF_QITEM, 0);
//    *pTaskType = TMQ_DELAYED_TASK__REPORT;
//    taosWriteQitem(tmq->delayedTask, pTaskType);
//    tsem_post(&tmq->rspSem);
//  }
//
//  taosReleaseRef(tmqMgmt.rsetId, refId);
//  taosMemoryFree(param);
//}
L
Liu Jicong 已提交
759

760
int32_t tmqHbCb(void* param, SDataBuf* pMsg, int32_t code) {
dengyihao's avatar
dengyihao 已提交
761 762 763 764
  if (pMsg) {
    taosMemoryFree(pMsg->pData);
    taosMemoryFree(pMsg->pEpSet);
  }
765 766 767 768
  return 0;
}

void tmqSendHbReq(void* param, void* tmrId) {
769
  int64_t refId = *(int64_t*)param;
770

X
Xiaoyu Wang 已提交
771
  tmq_t* tmq = taosAcquireRef(tmqMgmt.rsetId, refId);
772
  if (tmq == NULL) {
L
Liu Jicong 已提交
773
    taosMemoryFree(param);
774 775
    return;
  }
D
dapan1121 已提交
776 777 778 779

  SMqHbReq req = {0};
  req.consumerId = tmq->consumerId;
  req.epoch = tmq->epoch;
wmmhello's avatar
wmmhello 已提交
780
  taosRLockLatch(&tmq->lock);
wmmhello's avatar
wmmhello 已提交
781
//  if(tmq->needReportOffsetRows){
782 783 784 785 786 787 788 789 790 791 792 793
    req.topics = taosArrayInit(taosArrayGetSize(tmq->clientTopics), sizeof(TopicOffsetRows));
    for(int i = 0; i < taosArrayGetSize(tmq->clientTopics); i++){
      SMqClientTopic* pTopic = taosArrayGet(tmq->clientTopics, i);
      int32_t         numOfVgroups = taosArrayGetSize(pTopic->vgs);
      TopicOffsetRows* data = taosArrayReserve(req.topics, 1);
      strcpy(data->topicName, pTopic->topicName);
      data->offsetRows = taosArrayInit(numOfVgroups, sizeof(OffsetRows));
      for(int j = 0; j < numOfVgroups; j++){
        SMqClientVg* pVg = taosArrayGet(pTopic->vgs, j);
        OffsetRows* offRows = taosArrayReserve(data->offsetRows, 1);
        offRows->vgId = pVg->vgId;
        offRows->rows = pVg->numOfRows;
794
        offRows->offset = pVg->offsetInfo.beginOffset;
795 796
        char buf[TSDB_OFFSET_LEN] = {0};
        tFormatOffset(buf, TSDB_OFFSET_LEN, &offRows->offset);
wmmhello's avatar
wmmhello 已提交
797
        tscInfo("consumer:0x%" PRIx64 ",report offset: vgId:%d, offset:%s, rows:%"PRId64, tmq->consumerId, offRows->vgId, buf, offRows->rows);
798
      }
799
    }
wmmhello's avatar
wmmhello 已提交
800 801
//    tmq->needReportOffsetRows = false;
//  }
wmmhello's avatar
wmmhello 已提交
802
  taosRUnLockLatch(&tmq->lock);
D
dapan1121 已提交
803

L
Liu Jicong 已提交
804
  int32_t tlen = tSerializeSMqHbReq(NULL, 0, &req);
D
dapan1121 已提交
805 806
  if (tlen < 0) {
    tscError("tSerializeSMqHbReq failed");
807
    goto OVER;
D
dapan1121 已提交
808
  }
809

L
Liu Jicong 已提交
810
  void* pReq = taosMemoryCalloc(1, tlen);
D
dapan1121 已提交
811 812
  if (tlen < 0) {
    tscError("failed to malloc MqHbReq msg, size:%d", tlen);
813
    goto OVER;
D
dapan1121 已提交
814
  }
815

D
dapan1121 已提交
816 817 818
  if (tSerializeSMqHbReq(pReq, tlen, &req) < 0) {
    tscError("tSerializeSMqHbReq %d failed", tlen);
    taosMemoryFree(pReq);
819
    goto OVER;
D
dapan1121 已提交
820
  }
821 822 823 824

  SMsgSendInfo* sendInfo = taosMemoryCalloc(1, sizeof(SMsgSendInfo));
  if (sendInfo == NULL) {
    taosMemoryFree(pReq);
L
Liu Jicong 已提交
825
    goto OVER;
826
  }
827

828
  sendInfo->msgInfo = (SDataBuf){ .pData = pReq, .len = tlen, .handle = NULL };
829 830 831 832 833

  sendInfo->requestId = generateRequestId();
  sendInfo->requestObjRefId = 0;
  sendInfo->param = NULL;
  sendInfo->fp = tmqHbCb;
L
Liu Jicong 已提交
834
  sendInfo->msgType = TDMT_MND_TMQ_HB;
835 836 837 838 839 840 841

  SEpSet epSet = getEpSet_s(&tmq->pTscObj->pAppInfo->mgmtEp);

  int64_t transporterId = 0;
  asyncSendMsgToServer(tmq->pTscObj->pAppInfo->pTransporter, &epSet, &transporterId, sendInfo);

OVER:
842
  tDeatroySMqHbReq(&req);
843
  taosTmrReset(tmqSendHbReq, 1000, param, tmqMgmt.timer, &tmq->hbLiveTimer);
844
  taosReleaseRef(tmqMgmt.rsetId, refId);
845 846
}

847 848
static void defaultCommitCbFn(tmq_t* pTmq, int32_t code, void* param) {
  if (code != 0) {
849
    tscError("consumer:0x%" PRIx64 ", failed to commit offset, code:%s", pTmq->consumerId, tstrerror(code));
850 851 852
  }
}

853
int32_t tmqHandleAllDelayedTask(tmq_t* pTmq) {
L
Liu Jicong 已提交
854
  STaosQall* qall = taosAllocateQall();
855
  taosReadAllQitems(pTmq->delayedTask, qall);
L
Liu Jicong 已提交
856

857 858 859 860
  if (qall->numOfItems == 0) {
    taosFreeQall(qall);
    return TSDB_CODE_SUCCESS;
  }
861

X
Xiaoyu Wang 已提交
862
  tscDebug("consumer:0x%" PRIx64 " handle delayed %d tasks before poll data", pTmq->consumerId, qall->numOfItems);
863 864
  int8_t* pTaskType = NULL;
  taosGetQitem(qall, (void**)&pTaskType);
L
Liu Jicong 已提交
865

866
  while (pTaskType != NULL) {
867
    if (*pTaskType == TMQ_DELAYED_TASK__ASK_EP) {
868
      asyncAskEp(pTmq, addToQueueCallbackFn, NULL);
869 870

      int64_t* pRefId = taosMemoryMalloc(sizeof(int64_t));
871
      *pRefId = pTmq->refId;
872

X
Xiaoyu Wang 已提交
873
      tscDebug("consumer:0x%" PRIx64 " retrieve ep from mnode in 1s", pTmq->consumerId);
874
      taosTmrReset(tmqAssignAskEpTask, 1000, pRefId, tmqMgmt.timer, &pTmq->epTimer);
L
Liu Jicong 已提交
875
    } else if (*pTaskType == TMQ_DELAYED_TASK__COMMIT) {
X
Xiaoyu Wang 已提交
876
      tmq_commit_cb* pCallbackFn = pTmq->commitCb ? pTmq->commitCb : defaultCommitCbFn;
877 878

      asyncCommitAllOffsets(pTmq, pCallbackFn, pTmq->commitCbUserParam);
879
      int64_t* pRefId = taosMemoryMalloc(sizeof(int64_t));
880
      *pRefId = pTmq->refId;
881

882
      tscDebug("consumer:0x%" PRIx64 " next commit to vnode(s) in %.2fs", pTmq->consumerId,
X
Xiaoyu Wang 已提交
883
               pTmq->autoCommitInterval / 1000.0);
884
      taosTmrReset(tmqAssignDelayedCommitTask, pTmq->autoCommitInterval, pRefId, tmqMgmt.timer, &pTmq->commitTimer);
L
Liu Jicong 已提交
885 886
    } else if (*pTaskType == TMQ_DELAYED_TASK__REPORT) {
    }
887

L
Liu Jicong 已提交
888
    taosFreeQitem(pTaskType);
889
    taosGetQitem(qall, (void**)&pTaskType);
L
Liu Jicong 已提交
890
  }
891

L
Liu Jicong 已提交
892 893 894 895
  taosFreeQall(qall);
  return 0;
}

896
static void* tmqFreeRspWrapper(SMqRspWrapper* rspWrapper) {
L
Liu Jicong 已提交
897 898 899 900 901 902 903
  if (rspWrapper->tmqRspType == TMQ_MSG_TYPE__END_RSP) {
    // do nothing
  } else if (rspWrapper->tmqRspType == TMQ_MSG_TYPE__EP_RSP) {
    SMqAskEpRspWrapper* pEpRspWrapper = (SMqAskEpRspWrapper*)rspWrapper;
    tDeleteSMqAskEpRsp(&pEpRspWrapper->msg);
  } else if (rspWrapper->tmqRspType == TMQ_MSG_TYPE__POLL_RSP) {
    SMqPollRspWrapper* pRsp = (SMqPollRspWrapper*)rspWrapper;
904 905
    taosMemoryFreeClear(pRsp->pEpset);

L
Liu Jicong 已提交
906 907 908
    taosArrayDestroyP(pRsp->dataRsp.blockData, taosMemoryFree);
    taosArrayDestroy(pRsp->dataRsp.blockDataLen);
    taosArrayDestroyP(pRsp->dataRsp.blockTbName, taosMemoryFree);
909
    taosArrayDestroyP(pRsp->dataRsp.blockSchema, (FDelete)tDeleteSchemaWrapper);
L
Liu Jicong 已提交
910 911
  } else if (rspWrapper->tmqRspType == TMQ_MSG_TYPE__POLL_META_RSP) {
    SMqPollRspWrapper* pRsp = (SMqPollRspWrapper*)rspWrapper;
912 913
    taosMemoryFreeClear(pRsp->pEpset);

L
Liu Jicong 已提交
914 915 916
    taosMemoryFree(pRsp->metaRsp.metaRsp);
  } else if (rspWrapper->tmqRspType == TMQ_MSG_TYPE__TAOSX_RSP) {
    SMqPollRspWrapper* pRsp = (SMqPollRspWrapper*)rspWrapper;
917 918
    taosMemoryFreeClear(pRsp->pEpset);

L
Liu Jicong 已提交
919 920 921
    taosArrayDestroyP(pRsp->taosxRsp.blockData, taosMemoryFree);
    taosArrayDestroy(pRsp->taosxRsp.blockDataLen);
    taosArrayDestroyP(pRsp->taosxRsp.blockTbName, taosMemoryFree);
922
    taosArrayDestroyP(pRsp->taosxRsp.blockSchema, (FDelete)tDeleteSchemaWrapper);
L
Liu Jicong 已提交
923 924 925 926
    // taosx
    taosArrayDestroy(pRsp->taosxRsp.createTableLen);
    taosArrayDestroyP(pRsp->taosxRsp.createTableReq, taosMemoryFree);
  }
927 928

  return NULL;
L
Liu Jicong 已提交
929 930
}

L
Liu Jicong 已提交
931
void tmqClearUnhandleMsg(tmq_t* tmq) {
L
Liu Jicong 已提交
932
  SMqRspWrapper* rspWrapper = NULL;
L
Liu Jicong 已提交
933
  while (1) {
L
Liu Jicong 已提交
934 935 936 937 938
    taosGetQitem(tmq->qall, (void**)&rspWrapper);
    if (rspWrapper) {
      tmqFreeRspWrapper(rspWrapper);
      taosFreeQitem(rspWrapper);
    } else {
L
Liu Jicong 已提交
939
      break;
L
Liu Jicong 已提交
940
    }
L
Liu Jicong 已提交
941 942
  }

L
Liu Jicong 已提交
943
  rspWrapper = NULL;
L
Liu Jicong 已提交
944 945
  taosReadAllQitems(tmq->mqueue, tmq->qall);
  while (1) {
L
Liu Jicong 已提交
946 947 948 949 950
    taosGetQitem(tmq->qall, (void**)&rspWrapper);
    if (rspWrapper) {
      tmqFreeRspWrapper(rspWrapper);
      taosFreeQitem(rspWrapper);
    } else {
L
Liu Jicong 已提交
951
      break;
L
Liu Jicong 已提交
952
    }
L
Liu Jicong 已提交
953 954 955
  }
}

D
dapan1121 已提交
956
int32_t tmqSubscribeCb(void* param, SDataBuf* pMsg, int32_t code) {
L
Liu Jicong 已提交
957 958
  SMqSubscribeCbParam* pParam = (SMqSubscribeCbParam*)param;
  pParam->rspErr = code;
dengyihao's avatar
dengyihao 已提交
959 960

  taosMemoryFree(pMsg->pEpSet);
L
Liu Jicong 已提交
961 962 963
  tsem_post(&pParam->rspSem);
  return 0;
}
964

L
Liu Jicong 已提交
965
int32_t tmq_subscription(tmq_t* tmq, tmq_list_t** topics) {
966
  if(tmq == NULL) return TSDB_CODE_INVALID_PARA;
X
Xiaoyu Wang 已提交
967 968 969
  if (*topics == NULL) {
    *topics = tmq_list_new();
  }
wmmhello's avatar
wmmhello 已提交
970
  taosRLockLatch(&tmq->lock);
X
Xiaoyu Wang 已提交
971
  for (int i = 0; i < taosArrayGetSize(tmq->clientTopics); i++) {
L
Liu Jicong 已提交
972
    SMqClientTopic* topic = taosArrayGet(tmq->clientTopics, i);
L
Liu Jicong 已提交
973
    tmq_list_append(*topics, strchr(topic->topicName, '.') + 1);
X
Xiaoyu Wang 已提交
974
  }
wmmhello's avatar
wmmhello 已提交
975
  taosRUnLockLatch(&tmq->lock);
L
Liu Jicong 已提交
976
  return 0;
X
Xiaoyu Wang 已提交
977 978
}

L
Liu Jicong 已提交
979
int32_t tmq_unsubscribe(tmq_t* tmq) {
980
  if(tmq == NULL) return TSDB_CODE_INVALID_PARA;
981 982 983 984 985 986 987 988
  if (tmq->autoCommit) {
    int32_t rsp = tmq_commit_sync(tmq, NULL);
    if (rsp != 0) {
      return rsp;
    }
  }
  taosSsleep(2);  // sleep 2s for hb to send offset and rows to server

L
Liu Jicong 已提交
989 990
  int32_t     rsp;
  int32_t     retryCnt = 0;
L
Liu Jicong 已提交
991
  tmq_list_t* lst = tmq_list_new();
L
Liu Jicong 已提交
992 993 994 995 996 997 998 999 1000 1001
  while (1) {
    rsp = tmq_subscribe(tmq, lst);
    if (rsp != TSDB_CODE_MND_CONSUMER_NOT_READY || retryCnt > 5) {
      break;
    } else {
      retryCnt++;
      taosMsleep(500);
    }
  }

L
Liu Jicong 已提交
1002 1003
  tmq_list_destroy(lst);
  return rsp;
X
Xiaoyu Wang 已提交
1004 1005
}

1006 1007 1008 1009 1010 1011
static void freeClientVgImpl(void* param) {
  SMqClientTopic* pTopic = param;
  taosMemoryFreeClear(pTopic->schema.pSchema);
  taosArrayDestroy(pTopic->vgs);
}

1012
void tmqFreeImpl(void* handle) {
1013 1014
  tmq_t*  tmq = (tmq_t*)handle;
  int64_t id = tmq->consumerId;
L
Liu Jicong 已提交
1015

1016
  // TODO stop timer
L
Liu Jicong 已提交
1017 1018 1019 1020
  if (tmq->mqueue) {
    tmqClearUnhandleMsg(tmq);
    taosCloseQueue(tmq->mqueue);
  }
L
Liu Jicong 已提交
1021

H
Haojun Liao 已提交
1022 1023 1024 1025 1026
  if (tmq->delayedTask) {
    taosCloseQueue(tmq->delayedTask);
  }

  taosFreeQall(tmq->qall);
1027
  tsem_destroy(&tmq->rspSem);
L
Liu Jicong 已提交
1028

1029
  taosArrayDestroyEx(tmq->clientTopics, freeClientVgImpl);
1030 1031
  taos_close_internal(tmq->pTscObj);
  taosMemoryFree(tmq);
1032 1033

  tscDebug("consumer:0x%" PRIx64 " closed", id);
L
Liu Jicong 已提交
1034 1035
}

1036 1037 1038 1039 1040 1041 1042 1043 1044
static void tmqMgmtInit(void) {
  tmqInitRes = 0;
  tmqMgmt.timer = taosTmrInit(1000, 100, 360000, "TMQ");

  if (tmqMgmt.timer == NULL) {
    tmqInitRes = TSDB_CODE_OUT_OF_MEMORY;
  }

  tmqMgmt.rsetId = taosOpenRef(10000, tmqFreeImpl);
1045
  if (tmqMgmt.rsetId < 0) {
1046 1047 1048 1049
    tmqInitRes = terrno;
  }
}

L
Liu Jicong 已提交
1050
tmq_t* tmq_consumer_new(tmq_conf_t* conf, char* errstr, int32_t errstrLen) {
1051
  if(conf == NULL) return NULL;
1052 1053 1054 1055
  taosThreadOnce(&tmqInit, tmqMgmtInit);
  if (tmqInitRes != 0) {
    terrno = tmqInitRes;
    return NULL;
L
Liu Jicong 已提交
1056 1057
  }

L
Liu Jicong 已提交
1058 1059
  tmq_t* pTmq = taosMemoryCalloc(1, sizeof(tmq_t));
  if (pTmq == NULL) {
L
Liu Jicong 已提交
1060
    terrno = TSDB_CODE_OUT_OF_MEMORY;
1061
    tscError("failed to create consumer, groupId:%s, code:%s", conf->groupId, terrstr());
L
Liu Jicong 已提交
1062 1063
    return NULL;
  }
L
Liu Jicong 已提交
1064

L
Liu Jicong 已提交
1065 1066 1067
  const char* user = conf->user == NULL ? TSDB_DEFAULT_USER : conf->user;
  const char* pass = conf->pass == NULL ? TSDB_DEFAULT_PASS : conf->pass;

L
Liu Jicong 已提交
1068 1069 1070
  pTmq->clientTopics = taosArrayInit(0, sizeof(SMqClientTopic));
  pTmq->mqueue = taosOpenQueue();
  pTmq->delayedTask = taosOpenQueue();
H
Haojun Liao 已提交
1071
  pTmq->qall = taosAllocateQall();
L
Liu Jicong 已提交
1072

X
Xiaoyu Wang 已提交
1073 1074
  if (pTmq->clientTopics == NULL || pTmq->mqueue == NULL || pTmq->qall == NULL || pTmq->delayedTask == NULL ||
      conf->groupId[0] == 0) {
L
Liu Jicong 已提交
1075
    terrno = TSDB_CODE_OUT_OF_MEMORY;
X
Xiaoyu Wang 已提交
1076
    tscError("consumer:0x%" PRIx64 " setup failed since %s, groupId:%s", pTmq->consumerId, terrstr(), pTmq->groupId);
1077
    goto _failed;
L
Liu Jicong 已提交
1078
  }
L
Liu Jicong 已提交
1079

L
Liu Jicong 已提交
1080 1081
  // init status
  pTmq->status = TMQ_CONSUMER_STATUS__INIT;
L
Liu Jicong 已提交
1082 1083
  pTmq->pollCnt = 0;
  pTmq->epoch = 0;
wmmhello's avatar
wmmhello 已提交
1084
//  pTmq->needReportOffsetRows = true;
L
Liu Jicong 已提交
1085

L
Liu Jicong 已提交
1086 1087 1088
  // set conf
  strcpy(pTmq->clientId, conf->clientId);
  strcpy(pTmq->groupId, conf->groupId);
1089
  pTmq->withTbName = conf->withTbName;
L
Liu Jicong 已提交
1090
  pTmq->useSnapshot = conf->snapEnable;
L
Liu Jicong 已提交
1091
  pTmq->autoCommit = conf->autoCommit;
L
Liu Jicong 已提交
1092
  pTmq->autoCommitInterval = conf->autoCommitInterval;
L
Liu Jicong 已提交
1093 1094
  pTmq->commitCb = conf->commitCb;
  pTmq->commitCbUserParam = conf->commitCbUserParam;
L
Liu Jicong 已提交
1095
  pTmq->resetOffsetCfg = conf->resetOffset;
wmmhello's avatar
wmmhello 已提交
1096
  taosInitRWLatch(&pTmq->lock);
L
Liu Jicong 已提交
1097

1098 1099
  pTmq->hbBgEnable = conf->hbBgEnable;

L
Liu Jicong 已提交
1100
  // assign consumerId
L
Liu Jicong 已提交
1101
  pTmq->consumerId = tGenIdPI64();
X
Xiaoyu Wang 已提交
1102

L
Liu Jicong 已提交
1103 1104
  // init semaphore
  if (tsem_init(&pTmq->rspSem, 0, 0) != 0) {
1105
    tscError("consumer:0x %" PRIx64 " setup failed since %s, consumer group %s", pTmq->consumerId, terrstr(),
S
Shengliang Guan 已提交
1106
             pTmq->groupId);
1107
    goto _failed;
L
Liu Jicong 已提交
1108
  }
L
Liu Jicong 已提交
1109

L
Liu Jicong 已提交
1110 1111 1112
  // init connection
  pTmq->pTscObj = taos_connect_internal(conf->ip, user, pass, NULL, NULL, conf->port, CONN_TYPE__TMQ);
  if (pTmq->pTscObj == NULL) {
1113
    tscError("consumer:0x%" PRIx64 " setup failed since %s, groupId:%s", pTmq->consumerId, terrstr(), pTmq->groupId);
L
Liu Jicong 已提交
1114
    tsem_destroy(&pTmq->rspSem);
1115
    goto _failed;
L
Liu Jicong 已提交
1116
  }
L
Liu Jicong 已提交
1117

1118 1119
  pTmq->refId = taosAddRef(tmqMgmt.rsetId, pTmq);
  if (pTmq->refId < 0) {
1120
    goto _failed;
1121 1122
  }

1123
  if (pTmq->hbBgEnable) {
L
Liu Jicong 已提交
1124 1125
    int64_t* pRefId = taosMemoryMalloc(sizeof(int64_t));
    *pRefId = pTmq->refId;
1126
    pTmq->hbLiveTimer = taosTmrStart(tmqSendHbReq, 1000, pRefId, tmqMgmt.timer);
1127 1128
  }

1129
  char         buf[TSDB_OFFSET_LEN] = {0};
1130 1131
  STqOffsetVal offset = {.type = pTmq->resetOffsetCfg};
  tFormatOffset(buf, tListLen(buf), &offset);
X
Xiaoyu Wang 已提交
1132 1133 1134 1135
  tscInfo("consumer:0x%" PRIx64 " is setup, refId:%" PRId64
          ", groupId:%s, snapshot:%d, autoCommit:%d, commitInterval:%dms, offset:%s, backgroudHB:%d",
          pTmq->consumerId, pTmq->refId, pTmq->groupId, pTmq->useSnapshot, pTmq->autoCommit, pTmq->autoCommitInterval,
          buf, pTmq->hbBgEnable);
L
Liu Jicong 已提交
1136

1137
  return pTmq;
1138

1139 1140
_failed:
  tmqFreeImpl(pTmq);
L
Liu Jicong 已提交
1141
  return NULL;
1142 1143
}

L
Liu Jicong 已提交
1144
int32_t tmq_subscribe(tmq_t* tmq, const tmq_list_t* topic_list) {
1145
  if(tmq == NULL) return TSDB_CODE_INVALID_PARA;
1146
  const int32_t   MAX_RETRY_COUNT = 120 * 2;  // let's wait for 2 mins at most
L
Liu Jicong 已提交
1147 1148 1149
  const SArray*   container = &topic_list->container;
  int32_t         sz = taosArrayGetSize(container);
  void*           buf = NULL;
L
Liu Jicong 已提交
1150
  SMsgSendInfo*   sendInfo = NULL;
L
Liu Jicong 已提交
1151
  SCMSubscribeReq req = {0};
1152
  int32_t         code = 0;
1153

1154
  tscInfo("consumer:0x%" PRIx64 " cgroup:%s, subscribe %d topics", tmq->consumerId, tmq->groupId, sz);
L
Liu Jicong 已提交
1155

1156
  req.consumerId = tmq->consumerId;
L
Liu Jicong 已提交
1157
  tstrncpy(req.clientId, tmq->clientId, 256);
L
Liu Jicong 已提交
1158
  tstrncpy(req.cgroup, tmq->groupId, TSDB_CGROUP_LEN);
1159 1160
  req.topicNames = taosArrayInit(sz, sizeof(void*));

1161 1162 1163 1164
  if (req.topicNames == NULL) {
    code = TSDB_CODE_OUT_OF_MEMORY;
    goto FAIL;
  }
1165

1166 1167 1168 1169 1170
  req.withTbName = tmq->withTbName;
  req.autoCommit = tmq->autoCommit;
  req.autoCommitInterval = tmq->autoCommitInterval;
  req.resetOffsetCfg = tmq->resetOffsetCfg;

L
Liu Jicong 已提交
1171 1172
  for (int32_t i = 0; i < sz; i++) {
    char* topic = taosArrayGetP(container, i);
1173 1174

    SName name = {0};
L
Liu Jicong 已提交
1175 1176 1177 1178
    tNameSetDbName(&name, tmq->pTscObj->acctId, topic, strlen(topic));
    char* topicFName = taosMemoryCalloc(1, TSDB_TOPIC_FNAME_LEN);
    if (topicFName == NULL) {
      goto FAIL;
1179 1180
    }

1181
    tNameExtractFullName(&name, topicFName);
1182
    tscInfo("consumer:0x%" PRIx64 " subscribe topic:%s", tmq->consumerId, topicFName);
L
Liu Jicong 已提交
1183 1184

    taosArrayPush(req.topicNames, &topicFName);
1185 1186
  }

L
Liu Jicong 已提交
1187
  int32_t tlen = tSerializeSCMSubscribeReq(NULL, &req);
1188

L
Liu Jicong 已提交
1189
  buf = taosMemoryMalloc(tlen);
1190 1191 1192 1193
  if (buf == NULL) {
    code = TSDB_CODE_OUT_OF_MEMORY;
    goto FAIL;
  }
L
Liu Jicong 已提交
1194

1195 1196 1197
  void* abuf = buf;
  tSerializeSCMSubscribeReq(&abuf, &req);

L
Liu Jicong 已提交
1198
  sendInfo = taosMemoryCalloc(1, sizeof(SMsgSendInfo));
1199 1200 1201 1202
  if (sendInfo == NULL) {
    code = TSDB_CODE_OUT_OF_MEMORY;
    goto FAIL;
  }
1203

H
Haojun Liao 已提交
1204
  SMqSubscribeCbParam param = { .rspErr = 0};
1205
  if (tsem_init(&param.rspSem, 0, 0) != 0) {
wmmhello's avatar
wmmhello 已提交
1206
    code = TSDB_CODE_TSC_INTERNAL_ERROR;
1207 1208
    goto FAIL;
  }
L
Liu Jicong 已提交
1209

1210
  sendInfo->msgInfo = (SDataBuf){.pData = buf, .len = tlen, .handle = NULL};
1211

L
Liu Jicong 已提交
1212 1213
  sendInfo->requestId = generateRequestId();
  sendInfo->requestObjRefId = 0;
L
Liu Jicong 已提交
1214 1215
  sendInfo->param = &param;
  sendInfo->fp = tmqSubscribeCb;
L
Liu Jicong 已提交
1216
  sendInfo->msgType = TDMT_MND_TMQ_SUBSCRIBE;
L
Liu Jicong 已提交
1217

1218 1219 1220 1221 1222
  SEpSet epSet = getEpSet_s(&tmq->pTscObj->pAppInfo->mgmtEp);

  int64_t transporterId = 0;
  asyncSendMsgToServer(tmq->pTscObj->pAppInfo->pTransporter, &epSet, &transporterId, sendInfo);

L
Liu Jicong 已提交
1223 1224
  // avoid double free if msg is sent
  buf = NULL;
L
Liu Jicong 已提交
1225
  sendInfo = NULL;
L
Liu Jicong 已提交
1226

L
Liu Jicong 已提交
1227 1228
  tsem_wait(&param.rspSem);
  tsem_destroy(&param.rspSem);
1229

1230 1231 1232 1233
  if (param.rspErr != 0) {
    code = param.rspErr;
    goto FAIL;
  }
L
Liu Jicong 已提交
1234

L
Liu Jicong 已提交
1235
  int32_t retryCnt = 0;
1236
  while (TSDB_CODE_MND_CONSUMER_NOT_READY == doAskEp(tmq)) {
1237
    if (retryCnt++ > MAX_RETRY_COUNT) {
wmmhello's avatar
wmmhello 已提交
1238
      tscError("consumer:0x%" PRIx64 ", mnd not ready for subscribe, retry:%d in 500ms", tmq->consumerId, retryCnt);
1239
      code = TSDB_CODE_MND_CONSUMER_NOT_READY;
L
Liu Jicong 已提交
1240 1241
      goto FAIL;
    }
1242

1243
    tscInfo("consumer:0x%" PRIx64 ", mnd not ready for subscribe, retry:%d in 500ms", tmq->consumerId, retryCnt);
L
Liu Jicong 已提交
1244 1245
    taosMsleep(500);
  }
1246

1247 1248
  // init ep timer
  if (tmq->epTimer == NULL) {
1249 1250 1251
    int64_t* pRefId1 = taosMemoryMalloc(sizeof(int64_t));
    *pRefId1 = tmq->refId;
    tmq->epTimer = taosTmrStart(tmqAssignAskEpTask, 1000, pRefId1, tmqMgmt.timer);
1252
  }
L
Liu Jicong 已提交
1253 1254

  // init auto commit timer
1255
  if (tmq->autoCommit && tmq->commitTimer == NULL) {
1256 1257 1258
    int64_t* pRefId2 = taosMemoryMalloc(sizeof(int64_t));
    *pRefId2 = tmq->refId;
    tmq->commitTimer = taosTmrStart(tmqAssignDelayedCommitTask, tmq->autoCommitInterval, pRefId2, tmqMgmt.timer);
L
Liu Jicong 已提交
1259 1260
  }

L
Liu Jicong 已提交
1261
FAIL:
L
Liu Jicong 已提交
1262
  taosArrayDestroyP(req.topicNames, taosMemoryFree);
L
Liu Jicong 已提交
1263
  taosMemoryFree(buf);
L
Liu Jicong 已提交
1264
  taosMemoryFree(sendInfo);
L
Liu Jicong 已提交
1265

L
Liu Jicong 已提交
1266
  return code;
1267 1268
}

L
Liu Jicong 已提交
1269
void tmq_conf_set_auto_commit_cb(tmq_conf_t* conf, tmq_commit_cb* cb, void* param) {
1270
  if(conf == NULL) return;
1271
  conf->commitCb = cb;
L
Liu Jicong 已提交
1272
  conf->commitCbUserParam = param;
L
Liu Jicong 已提交
1273
}
1274

wmmhello's avatar
wmmhello 已提交
1275 1276 1277 1278 1279 1280 1281 1282 1283 1284 1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302
static SMqClientVg* getVgInfo(tmq_t* tmq, char* topicName, int32_t  vgId){
  int32_t topicNumCur = taosArrayGetSize(tmq->clientTopics);
  for(int i = 0; i < topicNumCur; i++){
    SMqClientTopic* pTopicCur = taosArrayGet(tmq->clientTopics, i);
    if(strcmp(pTopicCur->topicName, topicName) == 0){
      int32_t vgNumCur = taosArrayGetSize(pTopicCur->vgs);
      for (int32_t j = 0; j < vgNumCur; j++) {
        SMqClientVg* pVgCur = taosArrayGet(pTopicCur->vgs, j);
        if(pVgCur->vgId == vgId){
          return pVgCur;
        }
      }
    }
  }
  return NULL;
}

static SMqClientTopic* getTopicInfo(tmq_t* tmq, char* topicName){
  int32_t topicNumCur = taosArrayGetSize(tmq->clientTopics);
  for(int i = 0; i < topicNumCur; i++){
    SMqClientTopic* pTopicCur = taosArrayGet(tmq->clientTopics, i);
    if(strcmp(pTopicCur->topicName, topicName) == 0){
      return pTopicCur;
    }
  }
  return NULL;
}

D
dapan1121 已提交
1303
int32_t tmqPollCb(void* param, SDataBuf* pMsg, int32_t code) {
X
Xiaoyu Wang 已提交
1304
  SMqPollCbParam* pParam = (SMqPollCbParam*)param;
1305 1306

  int64_t         refId = pParam->refId;
wmmhello's avatar
wmmhello 已提交
1307 1308
//  SMqClientVg*    pVg = pParam->pVg;
//  SMqClientTopic* pTopic = pParam->pTopic;
1309

1310
  tmq_t* tmq = taosAcquireRef(tmqMgmt.rsetId, refId);
1311 1312 1313
  if (tmq == NULL) {
    taosMemoryFree(pParam);
    taosMemoryFree(pMsg->pData);
dengyihao's avatar
dengyihao 已提交
1314
    taosMemoryFree(pMsg->pEpSet);
1315 1316 1317 1318
    terrno = TSDB_CODE_TMQ_CONSUMER_CLOSED;
    return -1;
  }

H
Haojun Liao 已提交
1319 1320 1321 1322
  int32_t  epoch = pParam->epoch;
  int32_t  vgId = pParam->vgId;
  uint64_t requestId = pParam->requestId;

L
Liu Jicong 已提交
1323
  if (code != 0) {
L
Liu Jicong 已提交
1324
    if (pMsg->pData) taosMemoryFree(pMsg->pData);
dengyihao's avatar
dengyihao 已提交
1325 1326
    if (pMsg->pEpSet) taosMemoryFree(pMsg->pEpSet);

H
Haojun Liao 已提交
1327
    // in case of consumer mismatch, wait for 500ms and retry
L
Liu Jicong 已提交
1328
    if (code == TSDB_CODE_TMQ_CONSUMER_MISMATCH) {
1329
//      taosMsleep(500);
L
Liu Jicong 已提交
1330
      atomic_store_8(&tmq->status, TMQ_CONSUMER_STATUS__RECOVER);
X
Xiaoyu Wang 已提交
1331 1332
      tscDebug("consumer:0x%" PRIx64 " wait for the re-balance, wait for 500ms and set status to be RECOVER",
               tmq->consumerId);
H
Haojun Liao 已提交
1333
    } else if (code == TSDB_CODE_TQ_NO_COMMITTED_OFFSET) {
S
Shengliang Guan 已提交
1334
      SMqPollRspWrapper* pRspWrapper = taosAllocateQitem(sizeof(SMqPollRspWrapper), DEF_QITEM, 0);
L
Liu Jicong 已提交
1335
      if (pRspWrapper == NULL) {
H
Haojun Liao 已提交
1336 1337
        tscWarn("consumer:0x%" PRIx64 " msg from vgId:%d discarded, epoch %d since out of memory, reqId:0x%" PRIx64,
                tmq->consumerId, vgId, epoch, requestId);
L
Liu Jicong 已提交
1338 1339
        goto CREATE_MSG_FAIL;
      }
H
Haojun Liao 已提交
1340

L
Liu Jicong 已提交
1341 1342
      pRspWrapper->tmqRspType = TMQ_MSG_TYPE__END_RSP;
      taosWriteQitem(tmq->mqueue, pRspWrapper);
1343 1344
//    } else if (code == TSDB_CODE_WAL_LOG_NOT_EXIST) {  // poll data while insert
//      taosMsleep(5);
wmmhello's avatar
wmmhello 已提交
1345 1346 1347
    } else{
      tscError("consumer:0x%" PRIx64 " msg from vgId:%d discarded, epoch %d, since %s, reqId:0x%" PRIx64, tmq->consumerId,
               vgId, epoch, tstrerror(code), requestId);
L
Liu Jicong 已提交
1348
    }
H
Haojun Liao 已提交
1349

L
fix txn  
Liu Jicong 已提交
1350
    goto CREATE_MSG_FAIL;
L
Liu Jicong 已提交
1351 1352
  }

X
Xiaoyu Wang 已提交
1353
  int32_t msgEpoch = ((SMqRspHead*)pMsg->pData)->epoch;
1354 1355
  int32_t clientEpoch = atomic_load_32(&tmq->epoch);
  if (msgEpoch < clientEpoch) {
L
Liu Jicong 已提交
1356
    // do not write into queue since updating epoch reset
X
Xiaoyu Wang 已提交
1357 1358
    tscWarn("consumer:0x%" PRIx64
            " msg discard from vgId:%d since from earlier epoch, rsp epoch %d, current epoch %d, reqId:0x%" PRIx64,
1359
            tmq->consumerId, vgId, msgEpoch, clientEpoch, requestId);
H
Haojun Liao 已提交
1360

1361
    tsem_post(&tmq->rspSem);
1362 1363
    taosReleaseRef(tmqMgmt.rsetId, refId);

L
Liu Jicong 已提交
1364
    taosMemoryFree(pMsg->pData);
dengyihao's avatar
dengyihao 已提交
1365
    taosMemoryFree(pMsg->pEpSet);
wmmhello's avatar
wmmhello 已提交
1366 1367
    taosMemoryFree(pParam);

X
Xiaoyu Wang 已提交
1368 1369 1370
    return 0;
  }

1371
  if (msgEpoch != clientEpoch) {
H
Haojun Liao 已提交
1372
    tscWarn("consumer:0x%" PRIx64 " mismatch rsp from vgId:%d, epoch %d, current epoch %d, reqId:0x%" PRIx64,
1373
            tmq->consumerId, vgId, msgEpoch, clientEpoch, requestId);
X
Xiaoyu Wang 已提交
1374 1375
  }

L
Liu Jicong 已提交
1376 1377 1378
  // handle meta rsp
  int8_t rspType = ((SMqRspHead*)pMsg->pData)->mqMsgType;

S
Shengliang Guan 已提交
1379
  SMqPollRspWrapper* pRspWrapper = taosAllocateQitem(sizeof(SMqPollRspWrapper), DEF_QITEM, 0);
L
Liu Jicong 已提交
1380
  if (pRspWrapper == NULL) {
L
Liu Jicong 已提交
1381
    taosMemoryFree(pMsg->pData);
dengyihao's avatar
dengyihao 已提交
1382
    taosMemoryFree(pMsg->pEpSet);
X
Xiaoyu Wang 已提交
1383 1384
    tscWarn("consumer:0x%" PRIx64 " msg discard from vgId:%d, epoch %d since out of memory", tmq->consumerId, vgId,
            epoch);
L
fix txn  
Liu Jicong 已提交
1385
    goto CREATE_MSG_FAIL;
L
Liu Jicong 已提交
1386
  }
L
Liu Jicong 已提交
1387

L
Liu Jicong 已提交
1388
  pRspWrapper->tmqRspType = rspType;
wmmhello's avatar
wmmhello 已提交
1389 1390
//  pRspWrapper->vgHandle = pVg;
//  pRspWrapper->topicHandle = pTopic;
H
Haojun Liao 已提交
1391
  pRspWrapper->reqId = requestId;
1392
  pRspWrapper->pEpset = pMsg->pEpSet;
wmmhello's avatar
wmmhello 已提交
1393 1394
  pRspWrapper->vgId = vgId;
  strcpy(pRspWrapper->topicName, pParam->topicName);
L
Liu Jicong 已提交
1395

1396
  pMsg->pEpSet = NULL;
L
Liu Jicong 已提交
1397
  if (rspType == TMQ_MSG_TYPE__POLL_RSP) {
L
Liu Jicong 已提交
1398 1399
    SDecoder decoder;
    tDecoderInit(&decoder, POINTER_SHIFT(pMsg->pData, sizeof(SMqRspHead)), pMsg->len - sizeof(SMqRspHead));
1400
    tDecodeMqDataRsp(&decoder, &pRspWrapper->dataRsp);
wmmhello's avatar
wmmhello 已提交
1401
    tDecoderClear(&decoder);
L
Liu Jicong 已提交
1402
    memcpy(&pRspWrapper->dataRsp, pMsg->pData, sizeof(SMqRspHead));
L
Liu Jicong 已提交
1403

wmmhello's avatar
wmmhello 已提交
1404
    char buf[TSDB_OFFSET_LEN] = {0};
1405
    tFormatOffset(buf, TSDB_OFFSET_LEN, &pRspWrapper->dataRsp.rspOffset);
H
Haojun Liao 已提交
1406
    tscDebug("consumer:0x%" PRIx64 " recv poll rsp, vgId:%d, req ver:%" PRId64 ", rsp:%s type %d, reqId:0x%" PRIx64,
H
Haojun Liao 已提交
1407
             tmq->consumerId, vgId, pRspWrapper->dataRsp.reqOffset.version, buf, rspType, requestId);
L
Liu Jicong 已提交
1408
  } else if (rspType == TMQ_MSG_TYPE__POLL_META_RSP) {
1409 1410
    SDecoder decoder;
    tDecoderInit(&decoder, POINTER_SHIFT(pMsg->pData, sizeof(SMqRspHead)), pMsg->len - sizeof(SMqRspHead));
1411
    tDecodeMqMetaRsp(&decoder, &pRspWrapper->metaRsp);
1412
    tDecoderClear(&decoder);
L
Liu Jicong 已提交
1413
    memcpy(&pRspWrapper->metaRsp, pMsg->pData, sizeof(SMqRspHead));
L
Liu Jicong 已提交
1414 1415 1416 1417 1418 1419
  } else if (rspType == TMQ_MSG_TYPE__TAOSX_RSP) {
    SDecoder decoder;
    tDecoderInit(&decoder, POINTER_SHIFT(pMsg->pData, sizeof(SMqRspHead)), pMsg->len - sizeof(SMqRspHead));
    tDecodeSTaosxRsp(&decoder, &pRspWrapper->taosxRsp);
    tDecoderClear(&decoder);
    memcpy(&pRspWrapper->taosxRsp, pMsg->pData, sizeof(SMqRspHead));
X
Xiaoyu Wang 已提交
1420 1421
  } else {  // invalid rspType
    tscError("consumer:0x%" PRIx64 " invalid rsp msg received, type:%d ignored", tmq->consumerId, rspType);
L
Liu Jicong 已提交
1422
  }
L
Liu Jicong 已提交
1423

L
Liu Jicong 已提交
1424
  taosMemoryFree(pMsg->pData);
H
Haojun Liao 已提交
1425
  taosWriteQitem(tmq->mqueue, pRspWrapper);
L
Liu Jicong 已提交
1426

1427
  int32_t total = taosQueueItemSize(tmq->mqueue);
H
Haojun Liao 已提交
1428
  tscDebug("consumer:0x%" PRIx64 " put poll res into mqueue, type:%d, vgId:%d, total in queue:%d, reqId:0x%" PRIx64,
1429
           tmq->consumerId, rspType, vgId, total, requestId);
H
Haojun Liao 已提交
1430

1431
  tsem_post(&tmq->rspSem);
1432
  taosReleaseRef(tmqMgmt.rsetId, refId);
wmmhello's avatar
wmmhello 已提交
1433
  taosMemoryFree(pParam);
1434

L
Liu Jicong 已提交
1435
  return 0;
H
Haojun Liao 已提交
1436

L
fix txn  
Liu Jicong 已提交
1437
CREATE_MSG_FAIL:
L
Liu Jicong 已提交
1438
  if (epoch == tmq->epoch) {
wmmhello's avatar
wmmhello 已提交
1439 1440 1441 1442 1443
    taosWLockLatch(&tmq->lock);
    SMqClientVg* pVg = getVgInfo(tmq, pParam->topicName, vgId);
    if(pVg){
      atomic_store_32(&pVg->vgStatus, TMQ_VG_STATUS__IDLE);
    }
wmmhello's avatar
wmmhello 已提交
1444
    taosWUnLockLatch(&tmq->lock);
L
Liu Jicong 已提交
1445
  }
H
Haojun Liao 已提交
1446

1447
  tsem_post(&tmq->rspSem);
1448
  taosReleaseRef(tmqMgmt.rsetId, refId);
wmmhello's avatar
wmmhello 已提交
1449
  taosMemoryFree(pParam);
1450

L
Liu Jicong 已提交
1451
  return -1;
1452 1453
}

H
Haojun Liao 已提交
1454
typedef struct SVgroupSaveInfo {
wmmhello's avatar
wmmhello 已提交
1455 1456
  STqOffsetVal currentOffset;
  STqOffsetVal commitOffset;
1457
  STqOffsetVal seekOffset;
H
Haojun Liao 已提交
1458 1459 1460
  int64_t      numOfRows;
} SVgroupSaveInfo;

H
Haojun Liao 已提交
1461 1462 1463 1464 1465 1466
static void initClientTopicFromRsp(SMqClientTopic* pTopic, SMqSubTopicEp* pTopicEp, SHashObj* pVgOffsetHashMap,
                                   tmq_t* tmq) {
  pTopic->schema = pTopicEp->schema;
  pTopicEp->schema.nCols = 0;
  pTopicEp->schema.pSchema = NULL;

X
Xiaoyu Wang 已提交
1467
  char    vgKey[TSDB_TOPIC_FNAME_LEN + 22];
H
Haojun Liao 已提交
1468 1469 1470 1471 1472
  int32_t vgNumGet = taosArrayGetSize(pTopicEp->vgs);

  tstrncpy(pTopic->topicName, pTopicEp->topic, TSDB_TOPIC_FNAME_LEN);
  tstrncpy(pTopic->db, pTopicEp->db, TSDB_DB_FNAME_LEN);

1473
  tscInfo("consumer:0x%" PRIx64 ", update topic:%s, new numOfVgs:%d", tmq->consumerId, pTopic->topicName, vgNumGet);
H
Haojun Liao 已提交
1474 1475 1476 1477
  pTopic->vgs = taosArrayInit(vgNumGet, sizeof(SMqClientVg));

  for (int32_t j = 0; j < vgNumGet; j++) {
    SMqSubVgEp* pVgEp = taosArrayGet(pTopicEp->vgs, j);
H
Haojun Liao 已提交
1478 1479

    makeTopicVgroupKey(vgKey, pTopic->topicName, pVgEp->vgId);
H
Haojun Liao 已提交
1480
    SVgroupSaveInfo* pInfo = taosHashGet(pVgOffsetHashMap, vgKey, strlen(vgKey));
H
Haojun Liao 已提交
1481

wmmhello's avatar
wmmhello 已提交
1482 1483
    STqOffsetVal offsetNew = {0};
    offsetNew.type = tmq->resetOffsetCfg;
H
Haojun Liao 已提交
1484 1485 1486 1487 1488

    SMqClientVg clientVg = {
        .pollCnt = 0,
        .vgId = pVgEp->vgId,
        .epSet = pVgEp->epSet,
wmmhello's avatar
wmmhello 已提交
1489
        .vgStatus = TMQ_VG_STATUS__IDLE,
H
Haojun Liao 已提交
1490
        .vgSkipCnt = 0,
H
Haojun Liao 已提交
1491
        .emptyBlockReceiveTs = 0,
wmmhello's avatar
wmmhello 已提交
1492
        .numOfRows = pInfo ? pInfo->numOfRows : 0,
H
Haojun Liao 已提交
1493 1494
    };

1495
    clientVg.offsetInfo.endOffset = pInfo ? pInfo->currentOffset : offsetNew;
wmmhello's avatar
wmmhello 已提交
1496
    clientVg.offsetInfo.committedOffset = pInfo ? pInfo->commitOffset : offsetNew;
1497
    clientVg.offsetInfo.beginOffset = pInfo ? pInfo->seekOffset : offsetNew;
H
Haojun Liao 已提交
1498 1499
    clientVg.offsetInfo.walVerBegin = -1;
    clientVg.offsetInfo.walVerEnd = -1;
1500
    clientVg.seekUpdated = false;
1501
//    clientVg.receivedInfoFromVnode = false;
1502

H
Haojun Liao 已提交
1503 1504 1505 1506 1507 1508 1509 1510 1511 1512 1513 1514 1515
    taosArrayPush(pTopic->vgs, &clientVg);
  }
}

static void freeClientVgInfo(void* param) {
  SMqClientTopic* pTopic = param;
  if (pTopic->schema.nCols) {
    taosMemoryFreeClear(pTopic->schema.pSchema);
  }

  taosArrayDestroy(pTopic->vgs);
}

1516
static bool doUpdateLocalEp(tmq_t* tmq, int32_t epoch, const SMqAskEpRsp* pRsp) {
1517 1518 1519
  bool set = false;

  int32_t topicNumGet = taosArrayGetSize(pRsp->topics);
wmmhello's avatar
wmmhello 已提交
1520 1521 1522
  if (epoch <= tmq->epoch) {
    return false;
  }
1523 1524 1525 1526 1527 1528

  SArray* newTopics = taosArrayInit(topicNumGet, sizeof(SMqClientTopic));
  if (newTopics == NULL) {
    return false;
  }

H
Haojun Liao 已提交
1529 1530
  SHashObj* pVgOffsetHashMap = taosHashInit(64, MurmurHash3_32, false, HASH_NO_LOCK);
  if (pVgOffsetHashMap == NULL) {
1531 1532 1533
    taosArrayDestroy(newTopics);
    return false;
  }
1534

wmmhello's avatar
wmmhello 已提交
1535 1536 1537 1538 1539 1540
  taosWLockLatch(&tmq->lock);
  int32_t topicNumCur = taosArrayGetSize(tmq->clientTopics);

  char vgKey[TSDB_TOPIC_FNAME_LEN + 22];
  tscInfo("consumer:0x%" PRIx64 " update ep epoch from %d to epoch %d, incoming topics:%d, existed topics:%d",
          tmq->consumerId, tmq->epoch, epoch, topicNumGet, topicNumCur);
H
Haojun Liao 已提交
1541
  // todo extract method
1542 1543 1544 1545 1546
  for (int32_t i = 0; i < topicNumCur; i++) {
    // find old topic
    SMqClientTopic* pTopicCur = taosArrayGet(tmq->clientTopics, i);
    if (pTopicCur->vgs) {
      int32_t vgNumCur = taosArrayGetSize(pTopicCur->vgs);
1547
      tscInfo("consumer:0x%" PRIx64 ", current vg num: %d", tmq->consumerId, vgNumCur);
1548 1549
      for (int32_t j = 0; j < vgNumCur; j++) {
        SMqClientVg* pVgCur = taosArrayGet(pTopicCur->vgs, j);
H
Haojun Liao 已提交
1550 1551
        makeTopicVgroupKey(vgKey, pTopicCur->topicName, pVgCur->vgId);

wmmhello's avatar
wmmhello 已提交
1552
        char buf[TSDB_OFFSET_LEN] = {0};
1553
        tFormatOffset(buf, TSDB_OFFSET_LEN, &pVgCur->offsetInfo.endOffset);
1554
        tscInfo("consumer:0x%" PRIx64 ", epoch:%d vgId:%d vgKey:%s, offset:%s", tmq->consumerId, epoch, pVgCur->vgId,
X
Xiaoyu Wang 已提交
1555
                 vgKey, buf);
H
Haojun Liao 已提交
1556

1557
        SVgroupSaveInfo info = {.currentOffset = pVgCur->offsetInfo.endOffset, .seekOffset = pVgCur->offsetInfo.beginOffset, .commitOffset = pVgCur->offsetInfo.committedOffset, .numOfRows = pVgCur->numOfRows};
H
Haojun Liao 已提交
1558
        taosHashPut(pVgOffsetHashMap, vgKey, strlen(vgKey), &info, sizeof(SVgroupSaveInfo));
1559 1560 1561 1562 1563 1564 1565
      }
    }
  }

  for (int32_t i = 0; i < topicNumGet; i++) {
    SMqClientTopic topic = {0};
    SMqSubTopicEp* pTopicEp = taosArrayGet(pRsp->topics, i);
H
Haojun Liao 已提交
1566
    initClientTopicFromRsp(&topic, pTopicEp, pVgOffsetHashMap, tmq);
1567 1568
    taosArrayPush(newTopics, &topic);
  }
1569

H
Haojun Liao 已提交
1570 1571
  taosHashCleanup(pVgOffsetHashMap);

1572
  // destroy current buffered existed topics info
1573
  if (tmq->clientTopics) {
H
Haojun Liao 已提交
1574
    taosArrayDestroyEx(tmq->clientTopics, freeClientVgInfo);
X
Xiaoyu Wang 已提交
1575
  }
H
Haojun Liao 已提交
1576
  tmq->clientTopics = newTopics;
wmmhello's avatar
wmmhello 已提交
1577
  taosWUnLockLatch(&tmq->lock);
1578

X
Xiaoyu Wang 已提交
1579
  int8_t flag = (topicNumGet == 0) ? TMQ_CONSUMER_STATUS__NO_TOPIC : TMQ_CONSUMER_STATUS__READY;
H
Haojun Liao 已提交
1580
  atomic_store_8(&tmq->status, flag);
X
Xiaoyu Wang 已提交
1581
  atomic_store_32(&tmq->epoch, epoch);
H
Haojun Liao 已提交
1582

1583
  tscInfo("consumer:0x%" PRIx64 " update topic info completed", tmq->consumerId);
X
Xiaoyu Wang 已提交
1584 1585 1586
  return set;
}

1587
int32_t askEpCallbackFn(void* param, SDataBuf* pMsg, int32_t code) {
1588
  SMqAskEpCbParam* pParam = (SMqAskEpCbParam*)param;
1589 1590 1591
  tmq_t*           tmq = taosAcquireRef(tmqMgmt.rsetId, pParam->refId);

  if (tmq == NULL) {
1592
    terrno = TSDB_CODE_TMQ_CONSUMER_CLOSED;
wmmhello's avatar
wmmhello 已提交
1593
//    pParam->pUserFn(tmq, terrno, NULL, pParam->pParam);
1594

1595
    taosMemoryFree(pMsg->pData);
dengyihao's avatar
dengyihao 已提交
1596
    taosMemoryFree(pMsg->pEpSet);
1597 1598
    taosMemoryFree(pParam);
    return terrno;
1599 1600
  }

H
Haojun Liao 已提交
1601
  if (code != TSDB_CODE_SUCCESS) {
1602 1603 1604 1605 1606 1607 1608 1609 1610
    tscError("consumer:0x%" PRIx64 ", get topic endpoint error, code:%s", tmq->consumerId, tstrerror(code));
    pParam->pUserFn(tmq, code, NULL, pParam->pParam);

    taosReleaseRef(tmqMgmt.rsetId, pParam->refId);

    taosMemoryFree(pMsg->pData);
    taosMemoryFree(pMsg->pEpSet);
    taosMemoryFree(pParam);
    return code;
1611
  }
L
Liu Jicong 已提交
1612

L
Liu Jicong 已提交
1613
  // tmq's epoch is monotonically increase,
L
Liu Jicong 已提交
1614
  // so it's safe to discard any old epoch msg.
L
Liu Jicong 已提交
1615
  // Epoch will only increase when received newer epoch ep msg
L
Liu Jicong 已提交
1616 1617 1618
  SMqRspHead* head = pMsg->pData;
  int32_t     epoch = atomic_load_32(&tmq->epoch);
  if (head->epoch <= epoch) {
1619
    tscInfo("consumer:0x%" PRIx64 ", recv ep, msg epoch %d, current epoch %d, no need to update local ep",
1620
             tmq->consumerId, head->epoch, epoch);
1621

1622 1623 1624 1625 1626 1627 1628 1629
    if (tmq->status == TMQ_CONSUMER_STATUS__RECOVER) {
      SMqAskEpRsp rsp;
      tDecodeSMqAskEpRsp(POINTER_SHIFT(pMsg->pData, sizeof(SMqRspHead)), &rsp);
      int8_t flag = (taosArrayGetSize(rsp.topics) == 0) ? TMQ_CONSUMER_STATUS__NO_TOPIC : TMQ_CONSUMER_STATUS__READY;
      atomic_store_8(&tmq->status, flag);
      tDeleteSMqAskEpRsp(&rsp);
    }

X
Xiaoyu Wang 已提交
1630
  } else {
1631
    tscInfo("consumer:0x%" PRIx64 ", recv ep, msg epoch %d, current epoch %d, update local ep", tmq->consumerId,
1632
             head->epoch, epoch);
1633
  }
L
Liu Jicong 已提交
1634

1635
  pParam->pUserFn(tmq, code, pMsg, pParam->pParam);
1636 1637
  taosReleaseRef(tmqMgmt.rsetId, pParam->refId);

dengyihao's avatar
dengyihao 已提交
1638
  taosMemoryFree(pMsg->pEpSet);
L
Liu Jicong 已提交
1639
  taosMemoryFree(pMsg->pData);
1640
  taosMemoryFree(pParam);
L
Liu Jicong 已提交
1641
  return code;
1642 1643
}

L
Liu Jicong 已提交
1644
void tmqBuildConsumeReqImpl(SMqPollReq* pReq, tmq_t* tmq, int64_t timeout, SMqClientTopic* pTopic, SMqClientVg* pVg) {
L
Liu Jicong 已提交
1645 1646 1647 1648
  int32_t groupLen = strlen(tmq->groupId);
  memcpy(pReq->subKey, tmq->groupId, groupLen);
  pReq->subKey[groupLen] = TMQ_SEPARATOR;
  strcpy(pReq->subKey + groupLen + 1, pTopic->topicName);
1649

1650
  pReq->withTbName = tmq->withTbName;
L
Liu Jicong 已提交
1651
  pReq->consumerId = tmq->consumerId;
1652
  pReq->timeout = timeout;
X
Xiaoyu Wang 已提交
1653
  pReq->epoch = tmq->epoch;
1654
  pReq->reqOffset = pVg->offsetInfo.endOffset;
D
dapan1121 已提交
1655
  pReq->head.vgId = pVg->vgId;
1656 1657
  pReq->useSnapshot = tmq->useSnapshot;
  pReq->reqId = generateRequestId();
1658 1659
}

L
Liu Jicong 已提交
1660 1661
SMqMetaRspObj* tmqBuildMetaRspFromWrapper(SMqPollRspWrapper* pWrapper) {
  SMqMetaRspObj* pRspObj = taosMemoryCalloc(1, sizeof(SMqMetaRspObj));
L
Liu Jicong 已提交
1662
  pRspObj->resType = RES_TYPE__TMQ_META;
L
Liu Jicong 已提交
1663 1664 1665 1666 1667 1668 1669 1670
  tstrncpy(pRspObj->topic, pWrapper->topicHandle->topicName, TSDB_TOPIC_FNAME_LEN);
  tstrncpy(pRspObj->db, pWrapper->topicHandle->db, TSDB_DB_FNAME_LEN);
  pRspObj->vgId = pWrapper->vgHandle->vgId;

  memcpy(&pRspObj->metaRsp, &pWrapper->metaRsp, sizeof(SMqMetaRsp));
  return pRspObj;
}

1671
SMqRspObj* tmqBuildRspFromWrapper(SMqPollRspWrapper* pWrapper, SMqClientVg* pVg, int64_t* numOfRows) {
L
Liu Jicong 已提交
1672 1673
  SMqRspObj* pRspObj = taosMemoryCalloc(1, sizeof(SMqRspObj));
  pRspObj->resType = RES_TYPE__TMQ;
1674

1675
  (*numOfRows) = 0;
L
Liu Jicong 已提交
1676 1677
  tstrncpy(pRspObj->topic, pWrapper->topicHandle->topicName, TSDB_TOPIC_FNAME_LEN);
  tstrncpy(pRspObj->db, pWrapper->topicHandle->db, TSDB_DB_FNAME_LEN);
1678

L
Liu Jicong 已提交
1679
  pRspObj->vgId = pWrapper->vgHandle->vgId;
L
Liu Jicong 已提交
1680
  pRspObj->resIter = -1;
L
Liu Jicong 已提交
1681
  memcpy(&pRspObj->rsp, &pWrapper->dataRsp, sizeof(SMqDataRsp));
L
Liu Jicong 已提交
1682

L
Liu Jicong 已提交
1683 1684
  pRspObj->resInfo.totalRows = 0;
  pRspObj->resInfo.precision = TSDB_TIME_PRECISION_MILLI;
1685

L
Liu Jicong 已提交
1686
  if (!pWrapper->dataRsp.withSchema) {
L
Liu Jicong 已提交
1687 1688
    setResSchemaInfo(&pRspObj->resInfo, pWrapper->topicHandle->schema.pSchema, pWrapper->topicHandle->schema.nCols);
  }
L
Liu Jicong 已提交
1689

1690
  // extract the rows in this data packet
X
Xiaoyu Wang 已提交
1691
  for (int32_t i = 0; i < pRspObj->rsp.blockNum; ++i) {
1692
    SRetrieveTableRsp* pRetrieve = (SRetrieveTableRsp*)taosArrayGetP(pRspObj->rsp.blockData, i);
X
Xiaoyu Wang 已提交
1693
    int64_t            rows = htobe64(pRetrieve->numOfRows);
1694
    pVg->numOfRows += rows;
1695
    (*numOfRows) += rows;
1696 1697
  }

L
Liu Jicong 已提交
1698
  return pRspObj;
X
Xiaoyu Wang 已提交
1699 1700
}

1701
SMqTaosxRspObj* tmqBuildTaosxRspFromWrapper(SMqPollRspWrapper* pWrapper, SMqClientVg* pVg, int64_t* numOfRows) {
L
Liu Jicong 已提交
1702
  SMqTaosxRspObj* pRspObj = taosMemoryCalloc(1, sizeof(SMqTaosxRspObj));
1703
  pRspObj->resType = RES_TYPE__TMQ_METADATA;
L
Liu Jicong 已提交
1704 1705 1706 1707
  tstrncpy(pRspObj->topic, pWrapper->topicHandle->topicName, TSDB_TOPIC_FNAME_LEN);
  tstrncpy(pRspObj->db, pWrapper->topicHandle->db, TSDB_DB_FNAME_LEN);
  pRspObj->vgId = pWrapper->vgHandle->vgId;
  pRspObj->resIter = -1;
1708
  memcpy(&pRspObj->rsp, &pWrapper->taosxRsp, sizeof(STaosxRsp));
L
Liu Jicong 已提交
1709 1710 1711

  pRspObj->resInfo.totalRows = 0;
  pRspObj->resInfo.precision = TSDB_TIME_PRECISION_MILLI;
1712
  if (!pWrapper->taosxRsp.withSchema) {
L
Liu Jicong 已提交
1713 1714 1715
    setResSchemaInfo(&pRspObj->resInfo, pWrapper->topicHandle->schema.pSchema, pWrapper->topicHandle->schema.nCols);
  }

1716 1717 1718 1719 1720 1721 1722
  // extract the rows in this data packet
  for (int32_t i = 0; i < pRspObj->rsp.blockNum; ++i) {
    SRetrieveTableRsp* pRetrieve = (SRetrieveTableRsp*)taosArrayGetP(pRspObj->rsp.blockData, i);
    int64_t            rows = htobe64(pRetrieve->numOfRows);
    pVg->numOfRows += rows;
    (*numOfRows) += rows;
  }
L
Liu Jicong 已提交
1723 1724 1725
  return pRspObj;
}

1726 1727 1728 1729 1730 1731 1732 1733 1734 1735 1736 1737 1738 1739 1740 1741 1742 1743 1744 1745 1746 1747 1748 1749 1750 1751 1752 1753 1754 1755 1756 1757 1758
static int32_t handleErrorBeforePoll(SMqClientVg* pVg, tmq_t* pTmq) {
  atomic_store_32(&pVg->vgStatus, TMQ_VG_STATUS__IDLE);
  tsem_post(&pTmq->rspSem);
  return -1;
}

static int32_t doTmqPollImpl(tmq_t* pTmq, SMqClientTopic* pTopic, SMqClientVg* pVg, int64_t timeout) {
  SMqPollReq req = {0};
  tmqBuildConsumeReqImpl(&req, pTmq, timeout, pTopic, pVg);

  int32_t msgSize = tSerializeSMqPollReq(NULL, 0, &req);
  if (msgSize < 0) {
    return handleErrorBeforePoll(pVg, pTmq);
  }

  char* msg = taosMemoryCalloc(1, msgSize);
  if (NULL == msg) {
    return handleErrorBeforePoll(pVg, pTmq);
  }

  if (tSerializeSMqPollReq(msg, msgSize, &req) < 0) {
    taosMemoryFree(msg);
    return handleErrorBeforePoll(pVg, pTmq);
  }

  SMqPollCbParam* pParam = taosMemoryMalloc(sizeof(SMqPollCbParam));
  if (pParam == NULL) {
    taosMemoryFree(msg);
    return handleErrorBeforePoll(pVg, pTmq);
  }

  pParam->refId = pTmq->refId;
  pParam->epoch = pTmq->epoch;
wmmhello's avatar
wmmhello 已提交
1759 1760 1761
//  pParam->pVg = pVg;  // pVg may be released,fix it
//  pParam->pTopic = pTopic;
  strcpy(pParam->topicName, pTopic->topicName);
1762
  pParam->vgId = pVg->vgId;
H
Haojun Liao 已提交
1763
  pParam->requestId = req.reqId;
1764 1765 1766 1767 1768 1769 1770 1771

  SMsgSendInfo* sendInfo = taosMemoryCalloc(1, sizeof(SMsgSendInfo));
  if (sendInfo == NULL) {
    taosMemoryFree(pParam);
    taosMemoryFree(msg);
    return handleErrorBeforePoll(pVg, pTmq);
  }

H
Haojun Liao 已提交
1772
  sendInfo->msgInfo = (SDataBuf){.pData = msg, .len = msgSize, .handle = NULL};
1773 1774 1775 1776 1777 1778 1779
  sendInfo->requestId = req.reqId;
  sendInfo->requestObjRefId = 0;
  sendInfo->param = pParam;
  sendInfo->fp = tmqPollCb;
  sendInfo->msgType = TDMT_VND_TMQ_CONSUME;

  int64_t transporterId = 0;
wmmhello's avatar
wmmhello 已提交
1780
  char    offsetFormatBuf[TSDB_OFFSET_LEN] = {0};
1781
  tFormatOffset(offsetFormatBuf, tListLen(offsetFormatBuf), &pVg->offsetInfo.endOffset);
1782

X
Xiaoyu Wang 已提交
1783 1784
  tscDebug("consumer:0x%" PRIx64 " send poll to %s vgId:%d, epoch %d, req:%s, reqId:0x%" PRIx64, pTmq->consumerId,
           pTopic->topicName, pVg->vgId, pTmq->epoch, offsetFormatBuf, req.reqId);
1785 1786 1787
  asyncSendMsgToServer(pTmq->pTscObj->pAppInfo->pTransporter, &pVg->epSet, &transporterId, sendInfo);

  pVg->pollCnt++;
1788
  pVg->seekUpdated = false;   // reset this flag.
1789 1790 1791 1792 1793
  pTmq->pollCnt++;

  return TSDB_CODE_SUCCESS;
}

1794
// broadcast the poll request to all related vnodes
H
Haojun Liao 已提交
1795
static int32_t tmqPollImpl(tmq_t* tmq, int64_t timeout) {
1796 1797 1798
  if(atomic_load_8(&tmq->status) == TMQ_CONSUMER_STATUS__RECOVER){
    return 0;
  }
wmmhello's avatar
wmmhello 已提交
1799 1800 1801
  int32_t code = 0;

  taosWLockLatch(&tmq->lock);
1802
  int32_t numOfTopics = taosArrayGetSize(tmq->clientTopics);
X
Xiaoyu Wang 已提交
1803
  tscDebug("consumer:0x%" PRIx64 " start to poll data, numOfTopics:%d", tmq->consumerId, numOfTopics);
1804 1805

  for (int i = 0; i < numOfTopics; i++) {
X
Xiaoyu Wang 已提交
1806
    SMqClientTopic* pTopic = taosArrayGet(tmq->clientTopics, i);
X
Xiaoyu Wang 已提交
1807
    int32_t         numOfVg = taosArrayGetSize(pTopic->vgs);
1808 1809

    for (int j = 0; j < numOfVg; j++) {
X
Xiaoyu Wang 已提交
1810
      SMqClientVg* pVg = taosArrayGet(pTopic->vgs, j);
wmmhello's avatar
wmmhello 已提交
1811
      if (taosGetTimestampMs() - pVg->emptyBlockReceiveTs < EMPTY_BLOCK_POLL_IDLE_DURATION) {  // less than 10ms
1812
        tscTrace("consumer:0x%" PRIx64 " epoch %d, vgId:%d idle for 10ms before start next poll", tmq->consumerId,
X
Xiaoyu Wang 已提交
1813
                 tmq->epoch, pVg->vgId);
H
Haojun Liao 已提交
1814 1815 1816
        continue;
      }

1817
      int32_t vgStatus = atomic_val_compare_exchange_32(&pVg->vgStatus, TMQ_VG_STATUS__IDLE, TMQ_VG_STATUS__WAIT);
1818
      if (vgStatus == TMQ_VG_STATUS__WAIT) {
L
Liu Jicong 已提交
1819
        int32_t vgSkipCnt = atomic_add_fetch_32(&pVg->vgSkipCnt, 1);
1820
        tscTrace("consumer:0x%" PRIx64 " epoch %d wait poll-rsp, skip vgId:%d skip cnt %d", tmq->consumerId, tmq->epoch,
X
Xiaoyu Wang 已提交
1821
                 pVg->vgId, vgSkipCnt);
X
Xiaoyu Wang 已提交
1822 1823
        continue;
      }
1824

L
Liu Jicong 已提交
1825
      atomic_store_32(&pVg->vgSkipCnt, 0);
wmmhello's avatar
wmmhello 已提交
1826
      code = doTmqPollImpl(tmq, pTopic, pVg, timeout);
1827
      if (code != TSDB_CODE_SUCCESS) {
wmmhello's avatar
wmmhello 已提交
1828
        goto end;
D
dapan1121 已提交
1829
      }
X
Xiaoyu Wang 已提交
1830 1831
    }
  }
1832

wmmhello's avatar
wmmhello 已提交
1833 1834 1835 1836
end:
  taosWUnLockLatch(&tmq->lock);
  tscDebug("consumer:0x%" PRIx64 " end to poll data, code:%d", tmq->consumerId, code);
  return code;
X
Xiaoyu Wang 已提交
1837 1838
}

H
Haojun Liao 已提交
1839
static int32_t tmqHandleNoPollRsp(tmq_t* tmq, SMqRspWrapper* rspWrapper, bool* pReset) {
L
Liu Jicong 已提交
1840
  if (rspWrapper->tmqRspType == TMQ_MSG_TYPE__EP_RSP) {
L
fix  
Liu Jicong 已提交
1841
    /*printf("ep %d %d\n", rspMsg->head.epoch, tmq->epoch);*/
L
Liu Jicong 已提交
1842 1843
    if (rspWrapper->epoch > atomic_load_32(&tmq->epoch)) {
      SMqAskEpRspWrapper* pEpRspWrapper = (SMqAskEpRspWrapper*)rspWrapper;
L
Liu Jicong 已提交
1844
      SMqAskEpRsp*        rspMsg = &pEpRspWrapper->msg;
1845
      doUpdateLocalEp(tmq, rspWrapper->epoch, rspMsg);
L
temp  
Liu Jicong 已提交
1846
      /*tmqClearUnhandleMsg(tmq);*/
L
Liu Jicong 已提交
1847
      tDeleteSMqAskEpRsp(rspMsg);
X
Xiaoyu Wang 已提交
1848 1849
      *pReset = true;
    } else {
L
Liu Jicong 已提交
1850
      tmqFreeRspWrapper(rspWrapper);
X
Xiaoyu Wang 已提交
1851 1852 1853 1854 1855 1856 1857 1858
      *pReset = false;
    }
  } else {
    return -1;
  }
  return 0;
}

1859
static void updateVgInfo(SMqClientVg* pVg, STqOffsetVal* reqOffset, STqOffsetVal* rspOffset, int64_t sver, int64_t ever, int64_t consumerId){
wmmhello's avatar
wmmhello 已提交
1860 1861
  if (!pVg->seekUpdated) {
    tscDebug("consumer:0x%" PRIx64" local offset is update, since seekupdate not set", consumerId);
wmmhello's avatar
wmmhello 已提交
1862 1863
    if(reqOffset->type != 0) pVg->offsetInfo.beginOffset = *reqOffset;
    if(rspOffset->type != 0) pVg->offsetInfo.endOffset = *rspOffset;
wmmhello's avatar
wmmhello 已提交
1864 1865 1866 1867 1868 1869 1870 1871 1872
  } else {
    tscDebug("consumer:0x%" PRIx64" local offset is NOT update, since seekupdate is set", consumerId);
  }

  // update the status
  atomic_store_32(&pVg->vgStatus, TMQ_VG_STATUS__IDLE);

  // update the valid wal version range
  pVg->offsetInfo.walVerBegin = sver;
1873
  pVg->offsetInfo.walVerEnd = ever + 1;
1874
//  pVg->receivedInfoFromVnode = true;
wmmhello's avatar
wmmhello 已提交
1875 1876
}

H
Haojun Liao 已提交
1877
static void* tmqHandleAllRsp(tmq_t* tmq, int64_t timeout, bool pollIfReset) {
H
Haojun Liao 已提交
1878
  tscDebug("consumer:0x%" PRIx64 " start to handle the rsp, total:%d", tmq->consumerId, tmq->qall->numOfItems);
1879

X
Xiaoyu Wang 已提交
1880
  while (1) {
1881 1882
    SMqRspWrapper* pRspWrapper = NULL;
    taosGetQitem(tmq->qall, (void**)&pRspWrapper);
1883

1884
    if (pRspWrapper == NULL) {
L
Liu Jicong 已提交
1885
      taosReadAllQitems(tmq->mqueue, tmq->qall);
1886 1887
      taosGetQitem(tmq->qall, (void**)&pRspWrapper);
      if (pRspWrapper == NULL) {
L
Liu Jicong 已提交
1888 1889
        return NULL;
      }
X
Xiaoyu Wang 已提交
1890 1891
    }

X
Xiaoyu Wang 已提交
1892
    tscDebug("consumer:0x%" PRIx64 " handle rsp, type:%d", tmq->consumerId, pRspWrapper->tmqRspType);
H
Haojun Liao 已提交
1893

1894 1895
    if (pRspWrapper->tmqRspType == TMQ_MSG_TYPE__END_RSP) {
      taosFreeQitem(pRspWrapper);
L
Liu Jicong 已提交
1896
      terrno = TSDB_CODE_TQ_NO_COMMITTED_OFFSET;
H
Haojun Liao 已提交
1897
      tscError("consumer:0x%" PRIx64 " unexpected rsp from poll, code:%s", tmq->consumerId, tstrerror(terrno));
L
Liu Jicong 已提交
1898
      return NULL;
1899 1900
    } else if (pRspWrapper->tmqRspType == TMQ_MSG_TYPE__POLL_RSP) {
      SMqPollRspWrapper* pollRspWrapper = (SMqPollRspWrapper*)pRspWrapper;
H
Haojun Liao 已提交
1901

X
Xiaoyu Wang 已提交
1902
      int32_t     consumerEpoch = atomic_load_32(&tmq->epoch);
1903 1904 1905
      SMqDataRsp* pDataRsp = &pollRspWrapper->dataRsp;

      if (pDataRsp->head.epoch == consumerEpoch) {
wmmhello's avatar
wmmhello 已提交
1906
        taosWLockLatch(&tmq->lock);
wmmhello's avatar
wmmhello 已提交
1907 1908 1909 1910 1911 1912
        SMqClientVg* pVg = getVgInfo(tmq, pollRspWrapper->topicName, pollRspWrapper->vgId);
        pollRspWrapper->vgHandle = pVg;
        pollRspWrapper->topicHandle = getTopicInfo(tmq, pollRspWrapper->topicName);
        if(pollRspWrapper->vgHandle == NULL || pollRspWrapper->topicHandle == NULL){
          tscError("consumer:0x%" PRIx64 " get vg or topic error, topic:%s vgId:%d", tmq->consumerId,
                   pollRspWrapper->topicName, pollRspWrapper->vgId);
wmmhello's avatar
wmmhello 已提交
1913
          taosWUnLockLatch(&tmq->lock);
wmmhello's avatar
wmmhello 已提交
1914 1915
          return NULL;
        }
1916 1917 1918 1919
        // update the epset
        if (pollRspWrapper->pEpset != NULL) {
          SEp* pEp = GET_ACTIVE_EP(pollRspWrapper->pEpset);
          SEp* pOld = GET_ACTIVE_EP(&(pVg->epSet));
X
Xiaoyu Wang 已提交
1920 1921
          tscDebug("consumer:0x%" PRIx64 " update epset vgId:%d, ep:%s:%d, old ep:%s:%d", tmq->consumerId, pVg->vgId,
                   pEp->fqdn, pEp->port, pOld->fqdn, pOld->port);
1922 1923 1924
          pVg->epSet = *pollRspWrapper->pEpset;
        }

1925
        updateVgInfo(pVg, &pDataRsp->reqOffset, &pDataRsp->rspOffset, pDataRsp->head.walsver, pDataRsp->head.walever, tmq->consumerId);
1926

wmmhello's avatar
wmmhello 已提交
1927
        char buf[TSDB_OFFSET_LEN] = {0};
1928
        tFormatOffset(buf, TSDB_OFFSET_LEN, &pDataRsp->rspOffset);
1929
        if (pDataRsp->blockNum == 0) {
X
Xiaoyu Wang 已提交
1930
          tscDebug("consumer:0x%" PRIx64 " empty block received, vgId:%d, offset:%s, vg total:%" PRId64
wmmhello's avatar
wmmhello 已提交
1931
                   ", total:%" PRId64 ", reqId:0x%" PRIx64,
X
Xiaoyu Wang 已提交
1932
                   tmq->consumerId, pVg->vgId, buf, pVg->numOfRows, tmq->totalRows, pollRspWrapper->reqId);
1933
          pRspWrapper = tmqFreeRspWrapper(pRspWrapper);
1934
          pVg->emptyBlockReceiveTs = taosGetTimestampMs();
L
Liu Jicong 已提交
1935
          taosFreeQitem(pollRspWrapper);
1936
        } else {  // build rsp
X
Xiaoyu Wang 已提交
1937
          int64_t    numOfRows = 0;
1938
          SMqRspObj* pRsp = tmqBuildRspFromWrapper(pollRspWrapper, pVg, &numOfRows);
H
Haojun Liao 已提交
1939
          tmq->totalRows += numOfRows;
1940
          pVg->emptyBlockReceiveTs = 0;
H
Haojun Liao 已提交
1941
          tscDebug("consumer:0x%" PRIx64 " process poll rsp, vgId:%d, offset:%s, blocks:%d, rows:%" PRId64
wmmhello's avatar
wmmhello 已提交
1942
                   ", vg total:%" PRId64 ", total:%" PRId64 ", reqId:0x%" PRIx64,
H
Haojun Liao 已提交
1943
                   tmq->consumerId, pVg->vgId, buf, pDataRsp->blockNum, numOfRows, pVg->numOfRows, tmq->totalRows,
H
Haojun Liao 已提交
1944
                   pollRspWrapper->reqId);
1945
          taosFreeQitem(pollRspWrapper);
wmmhello's avatar
wmmhello 已提交
1946
          taosWUnLockLatch(&tmq->lock);
1947 1948
          return pRsp;
        }
wmmhello's avatar
wmmhello 已提交
1949
        taosWUnLockLatch(&tmq->lock);
X
Xiaoyu Wang 已提交
1950
      } else {
H
Haojun Liao 已提交
1951
        tscDebug("consumer:0x%" PRIx64 " vgId:%d msg discard since epoch mismatch: msg epoch %d, consumer epoch %d",
1952
                 tmq->consumerId, pollRspWrapper->vgId, pDataRsp->head.epoch, consumerEpoch);
1953
        pRspWrapper = tmqFreeRspWrapper(pRspWrapper);
L
Liu Jicong 已提交
1954 1955
        taosFreeQitem(pollRspWrapper);
      }
1956
    } else if (pRspWrapper->tmqRspType == TMQ_MSG_TYPE__POLL_META_RSP) {
1957
      // todo handle the wal range and epset for each vgroup
1958
      SMqPollRspWrapper* pollRspWrapper = (SMqPollRspWrapper*)pRspWrapper;
L
Liu Jicong 已提交
1959
      int32_t            consumerEpoch = atomic_load_32(&tmq->epoch);
1960 1961 1962

      tscDebug("consumer:0x%" PRIx64 " process meta rsp", tmq->consumerId);

L
Liu Jicong 已提交
1963
      if (pollRspWrapper->metaRsp.head.epoch == consumerEpoch) {
wmmhello's avatar
wmmhello 已提交
1964
        taosWLockLatch(&tmq->lock);
wmmhello's avatar
wmmhello 已提交
1965 1966 1967 1968 1969 1970
        SMqClientVg* pVg = getVgInfo(tmq, pollRspWrapper->topicName, pollRspWrapper->vgId);
        pollRspWrapper->vgHandle = pVg;
        pollRspWrapper->topicHandle = getTopicInfo(tmq, pollRspWrapper->topicName);
        if(pollRspWrapper->vgHandle == NULL || pollRspWrapper->topicHandle == NULL){
          tscError("consumer:0x%" PRIx64 " get vg or topic error, topic:%s vgId:%d", tmq->consumerId,
                   pollRspWrapper->topicName, pollRspWrapper->vgId);
wmmhello's avatar
wmmhello 已提交
1971
          taosWUnLockLatch(&tmq->lock);
wmmhello's avatar
wmmhello 已提交
1972 1973
          return NULL;
        }
H
Haojun Liao 已提交
1974

1975
        updateVgInfo(pVg, &pollRspWrapper->metaRsp.rspOffset, &pollRspWrapper->metaRsp.rspOffset, pollRspWrapper->metaRsp.head.walsver, pollRspWrapper->metaRsp.head.walever, tmq->consumerId);
L
Liu Jicong 已提交
1976
        // build rsp
L
Liu Jicong 已提交
1977
        SMqMetaRspObj* pRsp = tmqBuildMetaRspFromWrapper(pollRspWrapper);
L
Liu Jicong 已提交
1978
        taosFreeQitem(pollRspWrapper);
wmmhello's avatar
wmmhello 已提交
1979
        taosWUnLockLatch(&tmq->lock);
L
Liu Jicong 已提交
1980 1981
        return pRsp;
      } else {
H
Haojun Liao 已提交
1982
        tscDebug("consumer:0x%" PRIx64 " vgId:%d msg discard since epoch mismatch: msg epoch %d, consumer epoch %d",
1983
                 tmq->consumerId, pollRspWrapper->vgId, pollRspWrapper->metaRsp.head.epoch, consumerEpoch);
1984
        pRspWrapper = tmqFreeRspWrapper(pRspWrapper);
L
Liu Jicong 已提交
1985
        taosFreeQitem(pollRspWrapper);
X
Xiaoyu Wang 已提交
1986
      }
1987 1988
    } else if (pRspWrapper->tmqRspType == TMQ_MSG_TYPE__TAOSX_RSP) {
      SMqPollRspWrapper* pollRspWrapper = (SMqPollRspWrapper*)pRspWrapper;
X
Xiaoyu Wang 已提交
1989
      int32_t            consumerEpoch = atomic_load_32(&tmq->epoch);
H
Haojun Liao 已提交
1990

L
Liu Jicong 已提交
1991
      if (pollRspWrapper->taosxRsp.head.epoch == consumerEpoch) {
wmmhello's avatar
wmmhello 已提交
1992
        taosWLockLatch(&tmq->lock);
wmmhello's avatar
wmmhello 已提交
1993 1994 1995 1996 1997 1998
        SMqClientVg* pVg = getVgInfo(tmq, pollRspWrapper->topicName, pollRspWrapper->vgId);
        pollRspWrapper->vgHandle = pVg;
        pollRspWrapper->topicHandle = getTopicInfo(tmq, pollRspWrapper->topicName);
        if(pollRspWrapper->vgHandle == NULL || pollRspWrapper->topicHandle == NULL){
          tscError("consumer:0x%" PRIx64 " get vg or topic error, topic:%s vgId:%d", tmq->consumerId,
                   pollRspWrapper->topicName, pollRspWrapper->vgId);
wmmhello's avatar
wmmhello 已提交
1999
          taosWUnLockLatch(&tmq->lock);
wmmhello's avatar
wmmhello 已提交
2000 2001
          return NULL;
        }
H
Haojun Liao 已提交
2002

2003
        updateVgInfo(pVg, &pollRspWrapper->taosxRsp.reqOffset, &pollRspWrapper->taosxRsp.rspOffset, pollRspWrapper->taosxRsp.head.walsver, pollRspWrapper->taosxRsp.head.walever, tmq->consumerId);
H
Haojun Liao 已提交
2004

L
Liu Jicong 已提交
2005
        if (pollRspWrapper->taosxRsp.blockNum == 0) {
wmmhello's avatar
wmmhello 已提交
2006
          tscDebug("consumer:0x%" PRIx64 " taosx empty block received, vgId:%d, vg total:%" PRId64 ", reqId:0x%" PRIx64,
H
Haojun Liao 已提交
2007
                   tmq->consumerId, pVg->vgId, pVg->numOfRows, pollRspWrapper->reqId);
H
Haojun Liao 已提交
2008
          pVg->emptyBlockReceiveTs = taosGetTimestampMs();
2009
          pRspWrapper = tmqFreeRspWrapper(pRspWrapper);
H
Haojun Liao 已提交
2010
          taosFreeQitem(pollRspWrapper);
H
Haojun Liao 已提交
2011
        } else {
X
Xiaoyu Wang 已提交
2012
          pVg->emptyBlockReceiveTs = 0;  // reset the ts
wmmhello's avatar
wmmhello 已提交
2013 2014 2015 2016 2017 2018 2019 2020
          // build rsp
          void*   pRsp = NULL;
          int64_t numOfRows = 0;
          if (pollRspWrapper->taosxRsp.createTableNum == 0) {
            pRsp = tmqBuildRspFromWrapper(pollRspWrapper, pVg, &numOfRows);
          } else {
            pRsp = tmqBuildTaosxRspFromWrapper(pollRspWrapper, pVg, &numOfRows);
          }
2021

wmmhello's avatar
wmmhello 已提交
2022
          tmq->totalRows += numOfRows;
H
Haojun Liao 已提交
2023

wmmhello's avatar
wmmhello 已提交
2024
          char buf[TSDB_OFFSET_LEN] = {0};
2025
          tFormatOffset(buf, TSDB_OFFSET_LEN, &pVg->offsetInfo.endOffset);
wmmhello's avatar
wmmhello 已提交
2026 2027 2028 2029
          tscDebug("consumer:0x%" PRIx64 " process taosx poll rsp, vgId:%d, offset:%s, blocks:%d, rows:%" PRId64
                       ", vg total:%" PRId64 ", total:%" PRId64 ", reqId:0x%" PRIx64,
                   tmq->consumerId, pVg->vgId, buf, pollRspWrapper->dataRsp.blockNum, numOfRows, pVg->numOfRows,
                   tmq->totalRows, pollRspWrapper->reqId);
H
Haojun Liao 已提交
2030

wmmhello's avatar
wmmhello 已提交
2031 2032 2033 2034 2035
          taosFreeQitem(pollRspWrapper);
          taosWUnLockLatch(&tmq->lock);
          return pRsp;
        }
        taosWUnLockLatch(&tmq->lock);
L
Liu Jicong 已提交
2036
      } else {
H
Haojun Liao 已提交
2037
        tscDebug("consumer:0x%" PRIx64 " vgId:%d msg discard since epoch mismatch: msg epoch %d, consumer epoch %d",
2038
                 tmq->consumerId, pollRspWrapper->vgId, pollRspWrapper->taosxRsp.head.epoch, consumerEpoch);
2039
        pRspWrapper = tmqFreeRspWrapper(pRspWrapper);
L
Liu Jicong 已提交
2040 2041
        taosFreeQitem(pollRspWrapper);
      }
X
Xiaoyu Wang 已提交
2042
    } else {
H
Haojun Liao 已提交
2043 2044
      tscDebug("consumer:0x%" PRIx64 " not data msg received", tmq->consumerId);

X
Xiaoyu Wang 已提交
2045
      bool reset = false;
2046 2047
      tmqHandleNoPollRsp(tmq, pRspWrapper, &reset);
      taosFreeQitem(pRspWrapper);
X
Xiaoyu Wang 已提交
2048
      if (pollIfReset && reset) {
2049
        tscDebug("consumer:0x%" PRIx64 ", reset and repoll", tmq->consumerId);
2050
        tmqPollImpl(tmq, timeout);
X
Xiaoyu Wang 已提交
2051 2052 2053 2054 2055
      }
    }
  }
}

2056
TAOS_RES* tmq_consumer_poll(tmq_t* tmq, int64_t timeout) {
2057 2058
  if(tmq == NULL) return NULL;

L
Liu Jicong 已提交
2059 2060
  void*   rspObj;
  int64_t startTime = taosGetTimestampMs();
L
Liu Jicong 已提交
2061

2062
  tscInfo("consumer:0x%" PRIx64 " start to poll at %" PRId64 ", timeout:%" PRId64, tmq->consumerId, startTime,
X
Xiaoyu Wang 已提交
2063
           timeout);
L
Liu Jicong 已提交
2064

2065
  // in no topic status, delayed task also need to be processed
L
Liu Jicong 已提交
2066
  if (atomic_load_8(&tmq->status) == TMQ_CONSUMER_STATUS__INIT) {
2067
    tscInfo("consumer:0x%" PRIx64 " poll return since consumer is init", tmq->consumerId);
2068
    taosMsleep(500);  //     sleep for a while
2069 2070 2071
    return NULL;
  }

wmmhello's avatar
wmmhello 已提交
2072
  while (atomic_load_8(&tmq->status) == TMQ_CONSUMER_STATUS__RECOVER) {
L
Liu Jicong 已提交
2073
    int32_t retryCnt = 0;
2074
    while (TSDB_CODE_MND_CONSUMER_NOT_READY == doAskEp(tmq)) {
H
Haojun Liao 已提交
2075
      if (retryCnt++ > 40) {
L
Liu Jicong 已提交
2076 2077
        return NULL;
      }
2078

2079
      tscInfo("consumer:0x%" PRIx64 " not ready, retry:%d/40 in 500ms", tmq->consumerId, retryCnt);
L
Liu Jicong 已提交
2080 2081 2082 2083
      taosMsleep(500);
    }
  }

X
Xiaoyu Wang 已提交
2084
  while (1) {
L
Liu Jicong 已提交
2085
    tmqHandleAllDelayedTask(tmq);
2086

L
Liu Jicong 已提交
2087
    if (tmqPollImpl(tmq, timeout) < 0) {
2088
      tscError("consumer:0x%" PRIx64 " return due to poll error", tmq->consumerId);
L
Liu Jicong 已提交
2089
    }
L
Liu Jicong 已提交
2090

2091
    rspObj = tmqHandleAllRsp(tmq, timeout, false);
L
Liu Jicong 已提交
2092
    if (rspObj) {
2093
      tscDebug("consumer:0x%" PRIx64 " return rsp %p", tmq->consumerId, rspObj);
L
Liu Jicong 已提交
2094
      return (TAOS_RES*)rspObj;
L
Liu Jicong 已提交
2095
    } else if (terrno == TSDB_CODE_TQ_NO_COMMITTED_OFFSET) {
2096
      tscInfo("consumer:0x%" PRIx64 " return null since no committed offset", tmq->consumerId);
L
Liu Jicong 已提交
2097
      return NULL;
X
Xiaoyu Wang 已提交
2098
    }
2099

2100
    if (timeout >= 0) {
L
Liu Jicong 已提交
2101
      int64_t currentTime = taosGetTimestampMs();
2102 2103
      int64_t elapsedTime = currentTime - startTime;
      if (elapsedTime > timeout) {
2104
        tscInfo("consumer:0x%" PRIx64 " (epoch %d) timeout, no rsp, start time %" PRId64 ", current time %" PRId64,
L
Liu Jicong 已提交
2105
                 tmq->consumerId, tmq->epoch, startTime, currentTime);
X
Xiaoyu Wang 已提交
2106 2107
        return NULL;
      }
2108
      tsem_timewait(&tmq->rspSem, (timeout - elapsedTime));
L
Liu Jicong 已提交
2109 2110
    } else {
      // use tsem_timewait instead of tsem_wait to avoid unexpected stuck
L
Liu Jicong 已提交
2111
      tsem_timewait(&tmq->rspSem, 1000);
X
Xiaoyu Wang 已提交
2112 2113 2114 2115
    }
  }
}

wmmhello's avatar
wmmhello 已提交
2116 2117
static void displayConsumeStatistics(tmq_t* pTmq) {
  taosRLockLatch(&pTmq->lock);
2118 2119 2120 2121 2122 2123 2124 2125 2126 2127 2128 2129 2130
  int32_t numOfTopics = taosArrayGetSize(pTmq->clientTopics);
  tscDebug("consumer:0x%" PRIx64 " closing poll:%" PRId64 " rows:%" PRId64 " topics:%d, final epoch:%d",
           pTmq->consumerId, pTmq->pollCnt, pTmq->totalRows, numOfTopics, pTmq->epoch);

  tscDebug("consumer:0x%" PRIx64 " rows dist begin: ", pTmq->consumerId);
  for (int32_t i = 0; i < numOfTopics; ++i) {
    SMqClientTopic* pTopics = taosArrayGet(pTmq->clientTopics, i);

    tscDebug("consumer:0x%" PRIx64 " topic:%d", pTmq->consumerId, i);
    int32_t numOfVgs = taosArrayGetSize(pTopics->vgs);
    for (int32_t j = 0; j < numOfVgs; ++j) {
      SMqClientVg* pVg = taosArrayGet(pTopics->vgs, j);
      tscDebug("topic:%s, %d. vgId:%d rows:%" PRId64, pTopics->topicName, j, pVg->vgId, pVg->numOfRows);
2131
    }
2132
  }
wmmhello's avatar
wmmhello 已提交
2133
  taosRUnLockLatch(&pTmq->lock);
2134 2135
  tscDebug("consumer:0x%" PRIx64 " rows dist end", pTmq->consumerId);
}
2136

2137
int32_t tmq_consumer_close(tmq_t* tmq) {
2138 2139
  if(tmq == NULL) return TSDB_CODE_INVALID_PARA;

2140
  tscInfo("consumer:0x%" PRIx64 " start to close consumer, status:%d", tmq->consumerId, tmq->status);
2141
  displayConsumeStatistics(tmq);
2142

2143 2144 2145 2146 2147 2148
  if (tmq->status == TMQ_CONSUMER_STATUS__READY) {
    // if auto commit is set, commit before close consumer. Otherwise, do nothing.
    if (tmq->autoCommit) {
      int32_t rsp = tmq_commit_sync(tmq, NULL);
      if (rsp != 0) {
        return rsp;
2149 2150
      }
    }
2151
    taosSsleep(2);  // sleep 2s for hb to send offset and rows to server
2152

L
Liu Jicong 已提交
2153
    int32_t     retryCnt = 0;
2154
    tmq_list_t* lst = tmq_list_new();
L
Liu Jicong 已提交
2155
    while (1) {
2156
      int32_t rsp = tmq_subscribe(tmq, lst);
L
Liu Jicong 已提交
2157 2158 2159 2160 2161 2162 2163 2164
      if (rsp != TSDB_CODE_MND_CONSUMER_NOT_READY || retryCnt > 5) {
        break;
      } else {
        retryCnt++;
        taosMsleep(500);
      }
    }

2165
    tmq_list_destroy(lst);
2166
  } else {
2167
    tscInfo("consumer:0x%" PRIx64 " not in ready state, close it directly", tmq->consumerId);
L
Liu Jicong 已提交
2168
  }
H
Haojun Liao 已提交
2169

2170
  taosRemoveRef(tmqMgmt.rsetId, tmq->refId);
L
Liu Jicong 已提交
2171
  return 0;
2172
}
L
Liu Jicong 已提交
2173

L
Liu Jicong 已提交
2174 2175
const char* tmq_err2str(int32_t err) {
  if (err == 0) {
L
Liu Jicong 已提交
2176
    return "success";
L
Liu Jicong 已提交
2177
  } else if (err == -1) {
L
Liu Jicong 已提交
2178 2179 2180
    return "fail";
  } else {
    return tstrerror(err);
L
Liu Jicong 已提交
2181 2182
  }
}
L
Liu Jicong 已提交
2183

L
Liu Jicong 已提交
2184
tmq_res_t tmq_get_res_type(TAOS_RES* res) {
2185 2186 2187
  if (res == NULL){
    return TMQ_RES_INVALID;
  }
L
Liu Jicong 已提交
2188 2189 2190 2191
  if (TD_RES_TMQ(res)) {
    return TMQ_RES_DATA;
  } else if (TD_RES_TMQ_META(res)) {
    return TMQ_RES_TABLE_META;
2192 2193
  } else if (TD_RES_TMQ_METADATA(res)) {
    return TMQ_RES_METADATA;
L
Liu Jicong 已提交
2194 2195 2196 2197 2198
  } else {
    return TMQ_RES_INVALID;
  }
}

L
Liu Jicong 已提交
2199
const char* tmq_get_topic_name(TAOS_RES* res) {
2200 2201 2202
  if (res == NULL){
    return NULL;
  }
L
Liu Jicong 已提交
2203 2204
  if (TD_RES_TMQ(res)) {
    SMqRspObj* pRspObj = (SMqRspObj*)res;
L
Liu Jicong 已提交
2205
    return strchr(pRspObj->topic, '.') + 1;
L
Liu Jicong 已提交
2206 2207 2208
  } else if (TD_RES_TMQ_META(res)) {
    SMqMetaRspObj* pMetaRspObj = (SMqMetaRspObj*)res;
    return strchr(pMetaRspObj->topic, '.') + 1;
2209 2210 2211
  } else if (TD_RES_TMQ_METADATA(res)) {
    SMqTaosxRspObj* pRspObj = (SMqTaosxRspObj*)res;
    return strchr(pRspObj->topic, '.') + 1;
L
Liu Jicong 已提交
2212 2213 2214 2215 2216
  } else {
    return NULL;
  }
}

L
Liu Jicong 已提交
2217
const char* tmq_get_db_name(TAOS_RES* res) {
2218 2219 2220 2221
  if (res == NULL){
    return NULL;
  }

L
Liu Jicong 已提交
2222 2223 2224
  if (TD_RES_TMQ(res)) {
    SMqRspObj* pRspObj = (SMqRspObj*)res;
    return strchr(pRspObj->db, '.') + 1;
L
Liu Jicong 已提交
2225 2226 2227
  } else if (TD_RES_TMQ_META(res)) {
    SMqMetaRspObj* pMetaRspObj = (SMqMetaRspObj*)res;
    return strchr(pMetaRspObj->db, '.') + 1;
2228 2229 2230
  } else if (TD_RES_TMQ_METADATA(res)) {
    SMqTaosxRspObj* pRspObj = (SMqTaosxRspObj*)res;
    return strchr(pRspObj->db, '.') + 1;
L
Liu Jicong 已提交
2231 2232 2233 2234 2235
  } else {
    return NULL;
  }
}

L
Liu Jicong 已提交
2236
int32_t tmq_get_vgroup_id(TAOS_RES* res) {
2237 2238 2239
  if (res == NULL){
    return -1;
  }
L
Liu Jicong 已提交
2240 2241 2242
  if (TD_RES_TMQ(res)) {
    SMqRspObj* pRspObj = (SMqRspObj*)res;
    return pRspObj->vgId;
L
Liu Jicong 已提交
2243 2244 2245
  } else if (TD_RES_TMQ_META(res)) {
    SMqMetaRspObj* pMetaRspObj = (SMqMetaRspObj*)res;
    return pMetaRspObj->vgId;
2246
  } else if (TD_RES_TMQ_METADATA(res)) {
2247 2248
    SMqTaosxRspObj* pRspObj = (SMqTaosxRspObj*)res;
    return pRspObj->vgId;
L
Liu Jicong 已提交
2249 2250 2251 2252
  } else {
    return -1;
  }
}
L
Liu Jicong 已提交
2253

2254
int64_t tmq_get_vgroup_offset(TAOS_RES* res) {
2255 2256 2257
  if (res == NULL){
    return TSDB_CODE_INVALID_PARA;
  }
2258 2259
  if (TD_RES_TMQ(res)) {
    SMqRspObj* pRspObj = (SMqRspObj*) res;
2260
    STqOffsetVal* pOffset = &pRspObj->rsp.reqOffset;
2261
    if (pOffset->type == TMQ_OFFSET__LOG) {
2262
      return pRspObj->rsp.reqOffset.version;
2263 2264
    }else{
      tscError("invalid offset type:%d", pOffset->type);
2265 2266 2267 2268 2269 2270 2271 2272
    }
  } else if (TD_RES_TMQ_META(res)) {
    SMqMetaRspObj* pRspObj = (SMqMetaRspObj*)res;
    if (pRspObj->metaRsp.rspOffset.type == TMQ_OFFSET__LOG) {
      return pRspObj->metaRsp.rspOffset.version;
    }
  } else if (TD_RES_TMQ_METADATA(res)) {
    SMqTaosxRspObj* pRspObj = (SMqTaosxRspObj*) res;
2273 2274
    if (pRspObj->rsp.reqOffset.type == TMQ_OFFSET__LOG) {
      return pRspObj->rsp.reqOffset.version;
2275
    }
2276
  } else{
2277
    tscError("invalid tmq type:%d", *(int8_t*)res);
2278 2279 2280
  }

  // data from tsdb, no valid offset info
2281
  return TSDB_CODE_TMQ_SNAPSHOT_ERROR;
2282 2283
}

L
Liu Jicong 已提交
2284
const char* tmq_get_table_name(TAOS_RES* res) {
2285 2286 2287
  if (res == NULL){
    return NULL;
  }
L
Liu Jicong 已提交
2288 2289 2290 2291 2292 2293
  if (TD_RES_TMQ(res)) {
    SMqRspObj* pRspObj = (SMqRspObj*)res;
    if (!pRspObj->rsp.withTbName || pRspObj->rsp.blockTbName == NULL || pRspObj->resIter < 0 ||
        pRspObj->resIter >= pRspObj->rsp.blockNum) {
      return NULL;
    }
2294
    return (const char*)taosArrayGetP(pRspObj->rsp.blockTbName, pRspObj->resIter);
2295 2296
  } else if (TD_RES_TMQ_METADATA(res)) {
    SMqTaosxRspObj* pRspObj = (SMqTaosxRspObj*)res;
L
Liu Jicong 已提交
2297 2298 2299
    if (!pRspObj->rsp.withTbName || pRspObj->rsp.blockTbName == NULL || pRspObj->resIter < 0 ||
        pRspObj->resIter >= pRspObj->rsp.blockNum) {
      return NULL;
2300
    }
L
Liu Jicong 已提交
2301 2302
    return (const char*)taosArrayGetP(pRspObj->rsp.blockTbName, pRspObj->resIter);
  }
L
Liu Jicong 已提交
2303 2304
  return NULL;
}
2305

2306
void tmq_commit_async(tmq_t* tmq, const TAOS_RES* pRes, tmq_commit_cb* cb, void* param) {
2307 2308
  if (tmq == NULL) {
    tscError("invalid tmq handle, null");
2309 2310 2311
    if(cb != NULL) {
      cb(tmq, TSDB_CODE_INVALID_PARA, param);
    }
2312 2313
    return;
  }
2314 2315 2316
  if (pRes == NULL) {  // here needs to commit all offsets.
    asyncCommitAllOffsets(tmq, cb, param);
  } else {  // only commit one offset
2317
    asyncCommitFromResult(tmq, pRes, cb, param);
2318
  }
L
Liu Jicong 已提交
2319 2320
}

2321
static void commitCallBackFn(tmq_t *UNUSED_PARAM(tmq), int32_t code, void* param) {
2322 2323 2324
  SSyncCommitInfo* pInfo = (SSyncCommitInfo*) param;
  pInfo->code = code;
  tsem_post(&pInfo->sem);
2325
}
2326

2327
int32_t tmq_commit_sync(tmq_t* tmq, const TAOS_RES* pRes) {
2328 2329 2330 2331 2332
  if (tmq == NULL) {
    tscError("invalid tmq handle, null");
    return TSDB_CODE_INVALID_PARA;
  }

2333 2334 2335 2336 2337 2338 2339 2340
  int32_t code = 0;

  SSyncCommitInfo* pInfo = taosMemoryMalloc(sizeof(SSyncCommitInfo));
  tsem_init(&pInfo->sem, 0, 0);
  pInfo->code = 0;

  if (pRes == NULL) {
    asyncCommitAllOffsets(tmq, commitCallBackFn, pInfo);
H
Haojun Liao 已提交
2341
  } else {
2342
    asyncCommitFromResult(tmq, pRes, commitCallBackFn, pInfo);
2343 2344
  }

2345 2346
  tsem_wait(&pInfo->sem);
  code = pInfo->code;
H
Haojun Liao 已提交
2347 2348

  tsem_destroy(&pInfo->sem);
2349 2350
  taosMemoryFree(pInfo);

2351
  tscInfo("consumer:0x%" PRIx64 " sync res commit done, code:%s", tmq->consumerId, tstrerror(code));
2352 2353 2354
  return code;
}

2355
// wal range will be ok after calling tmq_get_topic_assignment or poll interface
2356 2357 2358 2359
static int32_t checkWalRange(SVgOffsetInfo* offset, int64_t value){
  if (offset->walVerBegin == -1 || offset->walVerEnd == -1) {
    tscError("Assignment or poll interface need to be called first");
    return TSDB_CODE_TMQ_NEED_INITIALIZED;
2360
  }
2361 2362 2363 2364 2365 2366 2367

  if (value != -1 && (value < offset->walVerBegin || value > offset->walVerEnd)) {
    tscError("invalid seek params, offset:%" PRId64 ", valid range:[%" PRId64 ", %" PRId64 "]", value, offset->walVerBegin, offset->walVerEnd);
    return TSDB_CODE_TMQ_VERSION_OUT_OF_RANGE;
  }

  return 0;
2368 2369 2370 2371 2372 2373 2374 2375 2376 2377 2378 2379 2380 2381
}

int32_t tmq_commit_offset_sync(tmq_t *tmq, const char *pTopicName, int32_t vgId, int64_t offset){
  if (tmq == NULL || pTopicName == NULL) {
    tscError("invalid tmq handle, null");
    return TSDB_CODE_INVALID_PARA;
  }

  int32_t accId = tmq->pTscObj->acctId;
  char tname[TSDB_TOPIC_FNAME_LEN] = {0};
  sprintf(tname, "%d.%s", accId, pTopicName);

  taosWLockLatch(&tmq->lock);
  SMqClientVg* pVg = NULL;
2382 2383
  int32_t code = getClientVg(tmq, tname, vgId, &pVg);
  if(code != 0){
2384
    taosWUnLockLatch(&tmq->lock);
2385
    return code;
2386 2387 2388
  }

  SVgOffsetInfo* pOffsetInfo = &pVg->offsetInfo;
2389 2390
  code = checkWalRange(pOffsetInfo, offset);
  if (code != 0) {
2391
    taosWUnLockLatch(&tmq->lock);
2392
    return code;
2393 2394 2395 2396 2397 2398 2399 2400 2401 2402 2403 2404 2405 2406
  }
  taosWUnLockLatch(&tmq->lock);

  STqOffsetVal offsetVal = {.type = TMQ_OFFSET__LOG, .version = offset};

  SSyncCommitInfo* pInfo = taosMemoryMalloc(sizeof(SSyncCommitInfo));
  if (pInfo == NULL) {
    tscError("consumer:0x%"PRIx64" failed to prepare seek operation", tmq->consumerId);
    return TSDB_CODE_OUT_OF_MEMORY;
  }

  tsem_init(&pInfo->sem, 0, 0);
  pInfo->code = 0;

2407 2408 2409 2410 2411
  code = asyncCommitOffset(tmq, tname, vgId, &offsetVal, commitCallBackFn, pInfo);
  if(code == 0){
    tsem_wait(&pInfo->sem);
    code = pInfo->code;
  }
2412 2413 2414 2415

  tsem_destroy(&pInfo->sem);
  taosMemoryFree(pInfo);

2416
  tscInfo("consumer:0x%" PRIx64 " sync send commit to vgId:%d, offset:%" PRId64" code:%s", tmq->consumerId, vgId, offset, tstrerror(code));
2417 2418 2419 2420

  return code;
}

2421 2422
void tmq_commit_offset_async(tmq_t *tmq, const char *pTopicName, int32_t vgId, int64_t offset, tmq_commit_cb *cb, void *param){
  int32_t code = 0;
2423 2424
  if (tmq == NULL || pTopicName == NULL) {
    tscError("invalid tmq handle, null");
2425 2426
    code = TSDB_CODE_INVALID_PARA;
    goto  end;
2427 2428 2429 2430 2431 2432 2433 2434
  }

  int32_t accId = tmq->pTscObj->acctId;
  char tname[TSDB_TOPIC_FNAME_LEN] = {0};
  sprintf(tname, "%d.%s", accId, pTopicName);

  taosWLockLatch(&tmq->lock);
  SMqClientVg* pVg = NULL;
2435
  code = getClientVg(tmq, tname, vgId, &pVg);
2436 2437
  if(code != 0){
    taosWUnLockLatch(&tmq->lock);
2438
    goto end;
2439 2440 2441 2442 2443 2444
  }

  SVgOffsetInfo* pOffsetInfo = &pVg->offsetInfo;
  code = checkWalRange(pOffsetInfo, offset);
  if (code != 0) {
    taosWUnLockLatch(&tmq->lock);
2445
    goto end;
2446 2447 2448 2449 2450 2451 2452
  }
  taosWUnLockLatch(&tmq->lock);

  STqOffsetVal offsetVal = {.type = TMQ_OFFSET__LOG, .version = offset};

  code = asyncCommitOffset(tmq, tname, vgId, &offsetVal, cb, param);

2453
  tscInfo("consumer:0x%" PRIx64 " async send commit to vgId:%d, offset:%" PRId64" code:%s", tmq->consumerId, vgId, offset, tstrerror(code));
2454

2455 2456 2457 2458
end:
  if(code != 0 && cb != NULL){
    cb(tmq, code, param);
  }
2459 2460
}

2461 2462 2463 2464 2465 2466 2467 2468 2469
void updateEpCallbackFn(tmq_t* pTmq, int32_t code, SDataBuf* pDataBuf, void* param) {
  SAskEpInfo* pInfo = param;
  pInfo->code = code;

  if (code == TSDB_CODE_SUCCESS) {
    SMqRspHead* head = pDataBuf->pData;

    SMqAskEpRsp rsp;
    tDecodeSMqAskEpRsp(POINTER_SHIFT(pDataBuf->pData, sizeof(SMqRspHead)), &rsp);
2470
    doUpdateLocalEp(pTmq, head->epoch, &rsp);
2471 2472 2473
    tDeleteSMqAskEpRsp(&rsp);
  }

H
Haojun Liao 已提交
2474
  tsem_post(&pInfo->sem);
2475 2476 2477 2478 2479 2480 2481 2482 2483 2484 2485 2486 2487 2488 2489 2490 2491 2492 2493 2494 2495 2496 2497 2498 2499 2500
}

void addToQueueCallbackFn(tmq_t* pTmq, int32_t code, SDataBuf* pDataBuf, void* param) {
  if (code != TSDB_CODE_SUCCESS) {
    terrno = code;
    return;
  }

  SMqAskEpRspWrapper* pWrapper = taosAllocateQitem(sizeof(SMqAskEpRspWrapper), DEF_QITEM, 0);
  if (pWrapper == NULL) {
    terrno = TSDB_CODE_OUT_OF_MEMORY;
    return;
  }

  SMqRspHead* head = pDataBuf->pData;

  pWrapper->tmqRspType = TMQ_MSG_TYPE__EP_RSP;
  pWrapper->epoch = head->epoch;
  memcpy(&pWrapper->msg, pDataBuf->pData, sizeof(SMqRspHead));
  tDecodeSMqAskEpRsp(POINTER_SHIFT(pDataBuf->pData, sizeof(SMqRspHead)), &pWrapper->msg);

  taosWriteQitem(pTmq->mqueue, pWrapper);
}

int32_t doAskEp(tmq_t* pTmq) {
  SAskEpInfo* pInfo = taosMemoryMalloc(sizeof(SAskEpInfo));
H
Haojun Liao 已提交
2501
  tsem_init(&pInfo->sem, 0, 0);
2502 2503

  asyncAskEp(pTmq, updateEpCallbackFn, pInfo);
H
Haojun Liao 已提交
2504
  tsem_wait(&pInfo->sem);
2505 2506

  int32_t code = pInfo->code;
H
Haojun Liao 已提交
2507
  tsem_destroy(&pInfo->sem);
2508 2509 2510 2511 2512
  taosMemoryFree(pInfo);
  return code;
}

void asyncAskEp(tmq_t* pTmq, __tmq_askep_fn_t askEpFn, void* param) {
2513
  SMqAskEpReq req = {0};
2514 2515 2516
  req.consumerId = pTmq->consumerId;
  req.epoch = pTmq->epoch;
  strcpy(req.cgroup, pTmq->groupId);
2517 2518 2519

  int32_t tlen = tSerializeSMqAskEpReq(NULL, 0, &req);
  if (tlen < 0) {
2520 2521 2522
    tscError("consumer:0x%" PRIx64 ", tSerializeSMqAskEpReq failed", pTmq->consumerId);
    askEpFn(pTmq, TSDB_CODE_INVALID_PARA, NULL, param);
    return;
2523 2524 2525 2526
  }

  void* pReq = taosMemoryCalloc(1, tlen);
  if (pReq == NULL) {
2527 2528 2529
    tscError("consumer:0x%" PRIx64 ", failed to malloc askEpReq msg, size:%d", pTmq->consumerId, tlen);
    askEpFn(pTmq, TSDB_CODE_OUT_OF_MEMORY, NULL, param);
    return;
2530 2531 2532
  }

  if (tSerializeSMqAskEpReq(pReq, tlen, &req) < 0) {
2533
    tscError("consumer:0x%" PRIx64 ", tSerializeSMqAskEpReq %d failed", pTmq->consumerId, tlen);
2534
    taosMemoryFree(pReq);
2535 2536 2537

    askEpFn(pTmq, TSDB_CODE_INVALID_PARA, NULL, param);
    return;
2538 2539 2540 2541
  }

  SMqAskEpCbParam* pParam = taosMemoryCalloc(1, sizeof(SMqAskEpCbParam));
  if (pParam == NULL) {
2542
    tscError("consumer:0x%" PRIx64 ", failed to malloc subscribe param", pTmq->consumerId);
2543
    taosMemoryFree(pReq);
2544 2545 2546

    askEpFn(pTmq, TSDB_CODE_OUT_OF_MEMORY, NULL, param);
    return;
2547 2548
  }

2549 2550 2551 2552
  pParam->refId = pTmq->refId;
  pParam->epoch = pTmq->epoch;
  pParam->pUserFn = askEpFn;
  pParam->pParam = param;
2553 2554 2555 2556 2557

  SMsgSendInfo* sendInfo = taosMemoryCalloc(1, sizeof(SMsgSendInfo));
  if (sendInfo == NULL) {
    taosMemoryFree(pParam);
    taosMemoryFree(pReq);
2558 2559
    askEpFn(pTmq, TSDB_CODE_OUT_OF_MEMORY, NULL, param);
    return;
2560 2561
  }

X
Xiaoyu Wang 已提交
2562
  sendInfo->msgInfo = (SDataBuf){.pData = pReq, .len = tlen, .handle = NULL};
2563 2564 2565 2566

  sendInfo->requestId = generateRequestId();
  sendInfo->requestObjRefId = 0;
  sendInfo->param = pParam;
2567
  sendInfo->fp = askEpCallbackFn;
2568 2569
  sendInfo->msgType = TDMT_MND_TMQ_ASK_EP;

2570
  SEpSet epSet = getEpSet_s(&pTmq->pTscObj->pAppInfo->mgmtEp);
2571
  tscInfo("consumer:0x%" PRIx64 " ask ep from mnode, reqId:0x%" PRIx64, pTmq->consumerId, sendInfo->requestId);
2572 2573

  int64_t transporterId = 0;
2574
  asyncSendMsgToServer(pTmq->pTscObj->pAppInfo->pTransporter, &epSet, &transporterId, sendInfo);
2575 2576 2577 2578 2579 2580 2581
}

int32_t makeTopicVgroupKey(char* dst, const char* topicName, int32_t vg) {
  return sprintf(dst, "%s:%d", topicName, vg);
}

int32_t tmqCommitDone(SMqCommitCbParamSet* pParamSet) {
2582 2583 2584
  int64_t refId = pParamSet->refId;

  tmq_t* tmq = taosAcquireRef(tmqMgmt.rsetId, refId);
2585 2586 2587 2588 2589 2590 2591
  if (tmq == NULL) {
    taosMemoryFree(pParamSet);
    terrno = TSDB_CODE_TMQ_CONSUMER_CLOSED;
    return -1;
  }

  // if no more waiting rsp
2592 2593 2594 2595
  if(pParamSet->callbackFn != NULL){
    pParamSet->callbackFn(tmq, pParamSet->code, pParamSet->userParam);
  }

2596
  taosMemoryFree(pParamSet);
wmmhello's avatar
wmmhello 已提交
2597
//  tmq->needReportOffsetRows = true;
2598 2599

  taosReleaseRef(tmqMgmt.rsetId, refId);
2600
  return 0;
2601 2602
}

2603
void commitRspCountDown(SMqCommitCbParamSet* pParamSet, int64_t consumerId, const char* pTopic, int32_t vgId) {
2604 2605
  int32_t waitingRspNum = atomic_sub_fetch_32(&pParamSet->waitingRspNum, 1);
  if (waitingRspNum == 0) {
2606
    tscInfo("consumer:0x%" PRIx64 " topic:%s vgId:%d all commit-rsp received, commit completed", consumerId, pTopic, vgId);
2607
    tmqCommitDone(pParamSet);
H
Haojun Liao 已提交
2608
  } else {
2609
    tscInfo("consumer:0x%" PRIx64 " topic:%s vgId:%d commit-rsp received, remain:%d", consumerId, pTopic, vgId, waitingRspNum);
2610 2611
  }
}
2612 2613 2614 2615 2616 2617 2618 2619 2620 2621 2622 2623 2624 2625 2626 2627 2628 2629 2630 2631 2632 2633

SReqResultInfo* tmqGetNextResInfo(TAOS_RES* res, bool convertUcs4) {
  SMqRspObj* pRspObj = (SMqRspObj*)res;
  pRspObj->resIter++;

  if (pRspObj->resIter < pRspObj->rsp.blockNum) {
    SRetrieveTableRsp* pRetrieve = (SRetrieveTableRsp*)taosArrayGetP(pRspObj->rsp.blockData, pRspObj->resIter);
    if (pRspObj->rsp.withSchema) {
      SSchemaWrapper* pSW = (SSchemaWrapper*)taosArrayGetP(pRspObj->rsp.blockSchema, pRspObj->resIter);
      setResSchemaInfo(&pRspObj->resInfo, pSW->pSchema, pSW->nCols);
      taosMemoryFreeClear(pRspObj->resInfo.row);
      taosMemoryFreeClear(pRspObj->resInfo.pCol);
      taosMemoryFreeClear(pRspObj->resInfo.length);
      taosMemoryFreeClear(pRspObj->resInfo.convertBuf);
      taosMemoryFreeClear(pRspObj->resInfo.convertJson);
    }

    setQueryResultFromRsp(&pRspObj->resInfo, pRetrieve, convertUcs4, false);
    return &pRspObj->resInfo;
  }

  return NULL;
H
Haojun Liao 已提交
2634 2635
}

2636 2637 2638 2639 2640 2641 2642 2643 2644 2645 2646 2647 2648 2649 2650 2651 2652 2653 2654
static int32_t tmqGetWalInfoCb(void* param, SDataBuf* pMsg, int32_t code) {
  SMqVgWalInfoParam* pParam = param;
  SMqVgCommon* pCommon = pParam->pCommon;

  int32_t total = atomic_add_fetch_32(&pCommon->numOfRsp, 1);
  if (code != TSDB_CODE_SUCCESS) {
    tscError("consumer:0x%" PRIx64 " failed to get the wal info from vgId:%d for topic:%s", pCommon->consumerId,
             pParam->vgId, pCommon->pTopicName);
    pCommon->code = code;
  } else {
    SMqDataRsp rsp;
    SDecoder decoder;
    tDecoderInit(&decoder, POINTER_SHIFT(pMsg->pData, sizeof(SMqRspHead)), pMsg->len - sizeof(SMqRspHead));
    tDecodeMqDataRsp(&decoder, &rsp);
    tDecoderClear(&decoder);

    SMqRspHead* pHead = pMsg->pData;

    tmq_topic_assignment assignment = {.begin = pHead->walsver,
2655
                                       .end = pHead->walever + 1,
2656
                                       .currentOffset = rsp.rspOffset.version,
2657
                                       .vgId = pParam->vgId};
2658 2659 2660 2661 2662 2663 2664 2665 2666 2667

    taosThreadMutexLock(&pCommon->mutex);
    taosArrayPush(pCommon->pList, &assignment);
    taosThreadMutexUnlock(&pCommon->mutex);
  }

  if (total == pParam->totalReq) {
    tsem_post(&pCommon->rsp);
  }

2668 2669
  taosMemoryFree(pMsg->pData);
  taosMemoryFree(pMsg->pEpSet);
2670 2671 2672 2673 2674
  taosMemoryFree(pParam);
  return 0;
}

static void destroyCommonInfo(SMqVgCommon* pCommon) {
wmmhello's avatar
wmmhello 已提交
2675 2676 2677
  if(pCommon == NULL){
    return;
  }
2678 2679 2680 2681 2682 2683 2684
  taosArrayDestroy(pCommon->pList);
  tsem_destroy(&pCommon->rsp);
  taosThreadMutexDestroy(&pCommon->mutex);
  taosMemoryFree(pCommon->pTopicName);
  taosMemoryFree(pCommon);
}

2685 2686 2687 2688 2689 2690 2691
static bool isInSnapshotMode(int8_t type, bool useSnapshot){
  if ((type < TMQ_OFFSET__LOG && useSnapshot) || type > TMQ_OFFSET__LOG) {
    return true;
  }
  return false;
}

2692 2693 2694 2695 2696 2697 2698 2699 2700 2701 2702 2703 2704 2705 2706 2707 2708 2709 2710 2711 2712 2713 2714 2715 2716 2717 2718 2719 2720 2721 2722 2723 2724 2725 2726 2727 2728 2729 2730 2731 2732 2733 2734 2735 2736 2737 2738 2739 2740 2741 2742 2743 2744 2745 2746 2747 2748 2749 2750 2751 2752 2753 2754 2755 2756 2757 2758 2759 2760 2761 2762 2763 2764 2765 2766 2767 2768 2769 2770 2771 2772 2773 2774 2775 2776 2777 2778 2779 2780 2781 2782 2783 2784 2785 2786 2787
static int32_t tmCommittedCb(void* param, SDataBuf* pMsg, int32_t code) {
  SMqCommittedParam* pParam = param;

  if (code != 0){
    goto end;
  }
  if (pMsg) {
    SDecoder decoder;
    tDecoderInit(&decoder, (uint8_t*)pMsg->pData, pMsg->len);
    if (tDecodeMqVgOffset(&decoder, &pParam->vgOffset) < 0) {
      code = TSDB_CODE_OUT_OF_MEMORY;
      goto end;
    }
    tDecoderClear(&decoder);
  }

  end:
  if(pMsg){
    taosMemoryFree(pMsg->pData);
    taosMemoryFree(pMsg->pEpSet);
  }
  pParam->code = code;
  tsem_post(&pParam->sem);
  return 0;
}

int64_t getCommittedFromServer(tmq_t *tmq, char* tname, int32_t vgId, SEpSet* epSet){
  int32_t code = 0;
  SMqVgOffset pOffset = {0};

  pOffset.consumerId = tmq->consumerId;

  int32_t groupLen = strlen(tmq->groupId);
  memcpy(pOffset.offset.subKey, tmq->groupId, groupLen);
  pOffset.offset.subKey[groupLen] = TMQ_SEPARATOR;
  strcpy(pOffset.offset.subKey + groupLen + 1, tname);

  int32_t len = 0;
  tEncodeSize(tEncodeMqVgOffset, &pOffset, len, code);
  if (code < 0) {
    return TSDB_CODE_INVALID_PARA;
  }

  void* buf = taosMemoryCalloc(1, sizeof(SMsgHead) + len);
  if (buf == NULL) {
    return TSDB_CODE_OUT_OF_MEMORY;
  }

  ((SMsgHead*)buf)->vgId = htonl(vgId);

  void* abuf = POINTER_SHIFT(buf, sizeof(SMsgHead));

  SEncoder encoder;
  tEncoderInit(&encoder, abuf, len);
  tEncodeMqVgOffset(&encoder, &pOffset);
  tEncoderClear(&encoder);

  SMsgSendInfo* sendInfo = taosMemoryCalloc(1, sizeof(SMsgSendInfo));
  if (sendInfo == NULL) {
    taosMemoryFree(buf);
    return TSDB_CODE_OUT_OF_MEMORY;
  }

  SMqCommittedParam* pParam = taosMemoryMalloc(sizeof(SMqCommittedParam));
  if (pParam == NULL) {
    taosMemoryFree(buf);
    taosMemoryFree(sendInfo);
    return TSDB_CODE_OUT_OF_MEMORY;
  }
  tsem_init(&pParam->sem, 0, 0);

  sendInfo->msgInfo = (SDataBuf){.pData = buf, .len = sizeof(SMsgHead) + len, .handle = NULL};
  sendInfo->requestId = generateRequestId();
  sendInfo->requestObjRefId = 0;
  sendInfo->param = pParam;
  sendInfo->fp = tmCommittedCb;
  sendInfo->msgType = TDMT_VND_TMQ_VG_COMMITTEDINFO;

  int64_t transporterId = 0;
  asyncSendMsgToServer(tmq->pTscObj->pAppInfo->pTransporter, epSet, &transporterId, sendInfo);

  tsem_wait(&pParam->sem);
  code = pParam->code;
  if(code == TSDB_CODE_SUCCESS){
    if(pParam->vgOffset.offset.val.type == TMQ_OFFSET__LOG){
      code = pParam->vgOffset.offset.val.version;
    }else{
      code = TSDB_CODE_TMQ_SNAPSHOT_ERROR;
    }
  }
  tsem_destroy(&pParam->sem);
  taosMemoryFree(pParam);

  return code;
}

2788
int64_t tmq_position(tmq_t *tmq, const char *pTopicName, int32_t vgId){
2789
  if (tmq == NULL || pTopicName == NULL) {
2790 2791 2792 2793 2794 2795 2796 2797 2798 2799 2800
    tscError("invalid tmq handle, null");
    return TSDB_CODE_INVALID_PARA;
  }

  int32_t accId = tmq->pTscObj->acctId;
  char tname[TSDB_TOPIC_FNAME_LEN] = {0};
  sprintf(tname, "%d.%s", accId, pTopicName);

  taosWLockLatch(&tmq->lock);

  SMqClientVg* pVg = NULL;
2801 2802
  int32_t code = getClientVg(tmq, tname, vgId, &pVg);
  if(code != 0){
2803
    taosWUnLockLatch(&tmq->lock);
2804
    return code;
2805 2806
  }

2807 2808
  SVgOffsetInfo* pOffsetInfo = &pVg->offsetInfo;
  int32_t type = pOffsetInfo->endOffset.type;
2809 2810 2811 2812 2813 2814
  if (isInSnapshotMode(type, tmq->useSnapshot)) {
    tscError("consumer:0x%" PRIx64 " offset type:%d not wal version, position error", tmq->consumerId, type);
    taosWUnLockLatch(&tmq->lock);
    return TSDB_CODE_TMQ_SNAPSHOT_ERROR;
  }

2815 2816
  code = checkWalRange(pOffsetInfo, -1);
  if (code != 0) {
2817
    taosWUnLockLatch(&tmq->lock);
2818
    return code;
2819
  }
2820 2821 2822 2823
  SEpSet epSet = pVg->epSet;
  int64_t begin = pVg->offsetInfo.walVerBegin;
  int64_t end = pVg->offsetInfo.walVerEnd;
  taosWUnLockLatch(&tmq->lock);
2824 2825 2826

  int64_t position = 0;
  if(type == TMQ_OFFSET__LOG){
2827 2828 2829 2830 2831 2832 2833 2834 2835 2836 2837 2838
    position = pOffsetInfo->endOffset.version;
  }else if(type == TMQ_OFFSET__RESET_EARLIEST || type == TMQ_OFFSET__RESET_LATEST){
    code = getCommittedFromServer(tmq, tname, vgId, &epSet);
    if(code == TSDB_CODE_TMQ_NO_COMMITTED){
      if(type == TMQ_OFFSET__RESET_EARLIEST){
        position = begin;
      } else if(type == TMQ_OFFSET__RESET_LATEST){
        position = end;
      }
    }else{
      position = code;
    }
2839 2840 2841 2842
  }else{
    tscError("consumer:0x%" PRIx64 " offset type:%d can not be reach here", tmq->consumerId, type);
  }

2843
  tscInfo("consumer:0x%" PRIx64 " tmq_position vgId:%d position:%" PRId64, tmq->consumerId, vgId, position);
2844 2845 2846
  return position;
}

2847 2848 2849 2850 2851 2852 2853 2854 2855 2856 2857 2858 2859 2860 2861 2862 2863 2864 2865 2866 2867 2868 2869 2870 2871 2872 2873 2874 2875 2876 2877 2878 2879 2880 2881 2882
int64_t tmq_committed(tmq_t *tmq, const char *pTopicName, int32_t vgId){
  if (tmq == NULL || pTopicName == NULL) {
    tscError("invalid tmq handle, null");
    return TSDB_CODE_INVALID_PARA;
  }

  int32_t accId = tmq->pTscObj->acctId;
  char tname[TSDB_TOPIC_FNAME_LEN] = {0};
  sprintf(tname, "%d.%s", accId, pTopicName);

  taosWLockLatch(&tmq->lock);

  SMqClientVg* pVg = NULL;
  int32_t code = getClientVg(tmq, tname, vgId, &pVg);
  if(code != 0){
    taosWUnLockLatch(&tmq->lock);
    return code;
  }

  SVgOffsetInfo* pOffsetInfo = &pVg->offsetInfo;
  if (isInSnapshotMode(pOffsetInfo->endOffset.type, tmq->useSnapshot)) {
    tscError("consumer:0x%" PRIx64 " offset type:%d not wal version, committed error", tmq->consumerId, pOffsetInfo->endOffset.type);
    taosWUnLockLatch(&tmq->lock);
    return TSDB_CODE_TMQ_SNAPSHOT_ERROR;
  }

  if (isInSnapshotMode(pOffsetInfo->committedOffset.type, tmq->useSnapshot)) {
    tscError("consumer:0x%" PRIx64 " offset type:%d not wal version, committed error", tmq->consumerId, pOffsetInfo->committedOffset.type);
    taosWUnLockLatch(&tmq->lock);
    return TSDB_CODE_TMQ_SNAPSHOT_ERROR;
  }

  int64_t committed = 0;
  if(pOffsetInfo->committedOffset.type == TMQ_OFFSET__LOG){
    committed = pOffsetInfo->committedOffset.version;
    taosWUnLockLatch(&tmq->lock);
2883
    goto end;
2884 2885 2886 2887
  }
  SEpSet epSet = pVg->epSet;
  taosWUnLockLatch(&tmq->lock);

2888 2889 2890 2891 2892
  committed = getCommittedFromServer(tmq, tname, vgId, &epSet);

end:
  tscInfo("consumer:0x%" PRIx64 " tmq_committed vgId:%d committed:%" PRId64, tmq->consumerId, vgId, committed);
  return committed;
2893 2894
}

H
Haojun Liao 已提交
2895
int32_t tmq_get_topic_assignment(tmq_t* tmq, const char* pTopicName, tmq_topic_assignment** assignment,
H
Haojun Liao 已提交
2896
                                 int32_t* numOfAssignment) {
2897 2898 2899 2900
  if(tmq == NULL || pTopicName == NULL || assignment == NULL || numOfAssignment == NULL){
    tscError("invalid tmq handle, null");
    return TSDB_CODE_INVALID_PARA;
  }
H
Haojun Liao 已提交
2901 2902
  *numOfAssignment = 0;
  *assignment = NULL;
wmmhello's avatar
wmmhello 已提交
2903
  SMqVgCommon* pCommon = NULL;
H
Haojun Liao 已提交
2904

2905
  int32_t accId = tmq->pTscObj->acctId;
2906
  char    tname[TSDB_TOPIC_FNAME_LEN] = {0};
2907
  sprintf(tname, "%d.%s", accId, pTopicName);
wmmhello's avatar
wmmhello 已提交
2908
  int32_t code = TSDB_CODE_SUCCESS;
2909

wmmhello's avatar
wmmhello 已提交
2910
  taosWLockLatch(&tmq->lock);
2911
  SMqClientTopic* pTopic = getTopicByName(tmq, tname);
H
Haojun Liao 已提交
2912
  if (pTopic == NULL) {
2913
    code = TSDB_CODE_TMQ_INVALID_TOPIC;
wmmhello's avatar
wmmhello 已提交
2914
    goto end;
H
Haojun Liao 已提交
2915 2916 2917 2918
  }

  // in case of snapshot is opened, no valid offset will return
  *numOfAssignment = taosArrayGetSize(pTopic->vgs);
2919 2920
  for (int32_t j = 0; j < (*numOfAssignment); ++j) {
    SMqClientVg* pClientVg = taosArrayGet(pTopic->vgs, j);
2921
    int32_t type = pClientVg->offsetInfo.beginOffset.type;
2922 2923
    if (isInSnapshotMode(type, tmq->useSnapshot)) {
      tscError("consumer:0x%" PRIx64 " offset type:%d not wal version, assignment not allowed", tmq->consumerId, type);
2924 2925 2926 2927
      code = TSDB_CODE_TMQ_SNAPSHOT_ERROR;
      goto end;
    }
  }
2928 2929 2930 2931 2932

  *assignment = taosMemoryCalloc(*numOfAssignment, sizeof(tmq_topic_assignment));
  if (*assignment == NULL) {
    tscError("consumer:0x%" PRIx64 " failed to malloc buffer, size:%" PRIzu, tmq->consumerId,
             (*numOfAssignment) * sizeof(tmq_topic_assignment));
wmmhello's avatar
wmmhello 已提交
2933 2934
    code = TSDB_CODE_OUT_OF_MEMORY;
    goto end;
2935 2936
  }

2937 2938
  bool needFetch = false;

H
Haojun Liao 已提交
2939 2940
  for (int32_t j = 0; j < (*numOfAssignment); ++j) {
    SMqClientVg* pClientVg = taosArrayGet(pTopic->vgs, j);
2941
    if (pClientVg->offsetInfo.beginOffset.type != TMQ_OFFSET__LOG) {
2942 2943 2944
      needFetch = true;
      break;
    }
H
Haojun Liao 已提交
2945 2946

    tmq_topic_assignment* pAssignment = &(*assignment)[j];
2947
    pAssignment->currentOffset = pClientVg->offsetInfo.beginOffset.version;
H
Haojun Liao 已提交
2948 2949
    pAssignment->begin = pClientVg->offsetInfo.walVerBegin;
    pAssignment->end = pClientVg->offsetInfo.walVerEnd;
2950
    pAssignment->vgId = pClientVg->vgId;
wmmhello's avatar
wmmhello 已提交
2951 2952
    tscInfo("consumer:0x%" PRIx64 " get assignment from local:%d->%" PRId64, tmq->consumerId,
            pAssignment->vgId, pAssignment->currentOffset);
H
Haojun Liao 已提交
2953 2954
  }

2955
  if (needFetch) {
wmmhello's avatar
wmmhello 已提交
2956
    pCommon = taosMemoryCalloc(1, sizeof(SMqVgCommon));
2957 2958
    if (pCommon == NULL) {
      terrno = TSDB_CODE_OUT_OF_MEMORY;
wmmhello's avatar
wmmhello 已提交
2959 2960
      code = terrno;
      goto end;
2961 2962 2963 2964 2965 2966 2967 2968 2969 2970 2971 2972 2973 2974
    }

    pCommon->pList= taosArrayInit(4, sizeof(tmq_topic_assignment));
    tsem_init(&pCommon->rsp, 0, 0);
    taosThreadMutexInit(&pCommon->mutex, 0);
    pCommon->pTopicName = taosStrdup(pTopic->topicName);
    pCommon->consumerId = tmq->consumerId;

    terrno = TSDB_CODE_OUT_OF_MEMORY;
    for (int32_t i = 0; i < (*numOfAssignment); ++i) {
      SMqClientVg* pClientVg = taosArrayGet(pTopic->vgs, i);

      SMqVgWalInfoParam* pParam = taosMemoryMalloc(sizeof(SMqVgWalInfoParam));
      if (pParam == NULL) {
wmmhello's avatar
wmmhello 已提交
2975 2976
        code = terrno;
        goto end;
2977 2978 2979 2980 2981 2982 2983 2984 2985
      }

      pParam->epoch = tmq->epoch;
      pParam->vgId = pClientVg->vgId;
      pParam->totalReq = *numOfAssignment;
      pParam->pCommon = pCommon;

      SMqPollReq req = {0};
      tmqBuildConsumeReqImpl(&req, tmq, 10, pTopic, pClientVg);
2986
      req.reqOffset = pClientVg->offsetInfo.beginOffset;
2987 2988 2989 2990

      int32_t msgSize = tSerializeSMqPollReq(NULL, 0, &req);
      if (msgSize < 0) {
        taosMemoryFree(pParam);
wmmhello's avatar
wmmhello 已提交
2991 2992
        code = terrno;
        goto end;
2993 2994 2995 2996 2997
      }

      char* msg = taosMemoryCalloc(1, msgSize);
      if (NULL == msg) {
        taosMemoryFree(pParam);
wmmhello's avatar
wmmhello 已提交
2998 2999
        code = terrno;
        goto end;
3000 3001 3002 3003 3004
      }

      if (tSerializeSMqPollReq(msg, msgSize, &req) < 0) {
        taosMemoryFree(msg);
        taosMemoryFree(pParam);
wmmhello's avatar
wmmhello 已提交
3005 3006
        code = terrno;
        goto end;
3007 3008 3009 3010 3011 3012
      }

      SMsgSendInfo* sendInfo = taosMemoryCalloc(1, sizeof(SMsgSendInfo));
      if (sendInfo == NULL) {
        taosMemoryFree(pParam);
        taosMemoryFree(msg);
wmmhello's avatar
wmmhello 已提交
3013 3014
        code = terrno;
        goto end;
3015 3016 3017 3018 3019 3020 3021 3022 3023 3024
      }

      sendInfo->msgInfo = (SDataBuf){.pData = msg, .len = msgSize, .handle = NULL};
      sendInfo->requestId = req.reqId;
      sendInfo->requestObjRefId = 0;
      sendInfo->param = pParam;
      sendInfo->fp = tmqGetWalInfoCb;
      sendInfo->msgType = TDMT_VND_TMQ_VG_WALINFO;

      int64_t transporterId = 0;
wmmhello's avatar
wmmhello 已提交
3025
      char    offsetFormatBuf[TSDB_OFFSET_LEN] = {0};
3026
      tFormatOffset(offsetFormatBuf, tListLen(offsetFormatBuf), &pClientVg->offsetInfo.beginOffset);
3027

3028
      tscInfo("consumer:0x%" PRIx64 " %s retrieve wal info vgId:%d, epoch %d, req:%s, reqId:0x%" PRIx64,
wmmhello's avatar
wmmhello 已提交
3029
              tmq->consumerId, pTopic->topicName, pClientVg->vgId, tmq->epoch, offsetFormatBuf, req.reqId);
3030 3031 3032 3033
      asyncSendMsgToServer(tmq->pTscObj->pAppInfo->pTransporter, &pClientVg->epSet, &transporterId, sendInfo);
    }

    tsem_wait(&pCommon->rsp);
wmmhello's avatar
wmmhello 已提交
3034
    code = pCommon->code;
3035 3036 3037

    terrno = code;
    if (code != TSDB_CODE_SUCCESS) {
wmmhello's avatar
wmmhello 已提交
3038
      goto end;
3039
    }
wmmhello's avatar
wmmhello 已提交
3040 3041 3042
    int32_t num = taosArrayGetSize(pCommon->pList);
    for(int32_t i = 0; i < num; ++i) {
      (*assignment)[i] = *(tmq_topic_assignment*)taosArrayGet(pCommon->pList, i);
3043
    }
wmmhello's avatar
wmmhello 已提交
3044
    *numOfAssignment = num;
3045

3046 3047 3048 3049 3050 3051 3052 3053 3054 3055
    for (int32_t j = 0; j < (*numOfAssignment); ++j) {
      tmq_topic_assignment* p = &(*assignment)[j];

      for(int32_t i = 0; i < taosArrayGetSize(pTopic->vgs); ++i) {
        SMqClientVg* pClientVg = taosArrayGet(pTopic->vgs, i);
        if (pClientVg->vgId != p->vgId) {
          continue;
        }

        SVgOffsetInfo* pOffsetInfo = &pClientVg->offsetInfo;
3056
        tscInfo("consumer:0x%" PRIx64 " %s vgId:%d offset is update to:%"PRId64, tmq->consumerId, pTopic->topicName, p->vgId, p->currentOffset);
3057 3058 3059 3060 3061

        pOffsetInfo->walVerBegin = p->begin;
        pOffsetInfo->walVerEnd = p->end;
      }
    }
wmmhello's avatar
wmmhello 已提交
3062
  }
3063

wmmhello's avatar
wmmhello 已提交
3064 3065 3066 3067 3068
end:
  if(code != TSDB_CODE_SUCCESS){
    taosMemoryFree(*assignment);
    *assignment = NULL;
    *numOfAssignment = 0;
3069
  }
wmmhello's avatar
wmmhello 已提交
3070 3071 3072
  destroyCommonInfo(pCommon);
  taosWUnLockLatch(&tmq->lock);
  return code;
H
Haojun Liao 已提交
3073 3074
}

T
t_max 已提交
3075 3076 3077 3078 3079 3080 3081 3082
void tmq_free_assignment(tmq_topic_assignment* pAssignment) {
    if (pAssignment == NULL) {
        return;
    }

    taosMemoryFree(pAssignment);
}

3083 3084 3085 3086 3087 3088 3089 3090 3091 3092 3093
static int32_t tmqSeekCb(void* param, SDataBuf* pMsg, int32_t code) {
  if (pMsg) {
    taosMemoryFree(pMsg->pData);
    taosMemoryFree(pMsg->pEpSet);
  }
  SMqSeekParam* pParam = param;
  pParam->code = code;
  tsem_post(&pParam->sem);
  return 0;
}

3094
// seek interface have to send msg to server to cancel push handle if needed, because consumer may be in wait status if there is no data to poll
3095
int32_t tmq_offset_seek(tmq_t* tmq, const char* pTopicName, int32_t vgId, int64_t offset) {
3096
  if (tmq == NULL || pTopicName == NULL) {
H
Haojun Liao 已提交
3097
    tscError("invalid tmq handle, null");
H
Haojun Liao 已提交
3098 3099 3100
    return TSDB_CODE_INVALID_PARA;
  }

3101
  int32_t accId = tmq->pTscObj->acctId;
3102
  char tname[TSDB_TOPIC_FNAME_LEN] = {0};
3103 3104
  sprintf(tname, "%d.%s", accId, pTopicName);

wmmhello's avatar
wmmhello 已提交
3105
  taosWLockLatch(&tmq->lock);
H
Haojun Liao 已提交
3106 3107

  SMqClientVg* pVg = NULL;
3108 3109
  int32_t code = getClientVg(tmq, tname, vgId, &pVg);
  if(code != 0){
wmmhello's avatar
wmmhello 已提交
3110
    taosWUnLockLatch(&tmq->lock);
3111
    return code;
H
Haojun Liao 已提交
3112 3113 3114 3115
  }

  SVgOffsetInfo* pOffsetInfo = &pVg->offsetInfo;

3116
  int32_t type = pOffsetInfo->endOffset.type;
3117
  if (isInSnapshotMode(type, tmq->useSnapshot)) {
wmmhello's avatar
wmmhello 已提交
3118 3119
    tscError("consumer:0x%" PRIx64 " offset type:%d not wal version, seek not allowed", tmq->consumerId, type);
    taosWUnLockLatch(&tmq->lock);
3120
    return TSDB_CODE_TMQ_SNAPSHOT_ERROR;
wmmhello's avatar
wmmhello 已提交
3121 3122
  }

3123
  code = checkWalRange(pOffsetInfo, offset);
3124
  if (code != 0) {
wmmhello's avatar
wmmhello 已提交
3125
    taosWUnLockLatch(&tmq->lock);
3126
    return code;
wmmhello's avatar
wmmhello 已提交
3127
  }
H
Haojun Liao 已提交
3128

3129
  tscInfo("consumer:0x%" PRIx64 " seek to %" PRId64 " on vgId:%d", tmq->consumerId, offset, vgId);
H
Haojun Liao 已提交
3130
  // update the offset, and then commit to vnode
3131 3132 3133
  pOffsetInfo->endOffset.type = TMQ_OFFSET__LOG;
  pOffsetInfo->endOffset.version = offset;
  pOffsetInfo->beginOffset = pOffsetInfo->endOffset;
wmmhello's avatar
wmmhello 已提交
3134
  pVg->seekUpdated = true;
3135 3136
  SEpSet epSet = pVg->epSet;
  taosWUnLockLatch(&tmq->lock);
3137 3138

  SMqSeekReq req = {0};
3139 3140
  snprintf(req.subKey, TSDB_SUBSCRIBE_KEY_LEN, "%s:%s", tmq->groupId, tname);
  req.head.vgId = vgId;
3141 3142 3143 3144 3145 3146 3147 3148 3149 3150 3151 3152 3153 3154 3155 3156 3157 3158 3159 3160 3161 3162 3163 3164 3165 3166 3167 3168 3169 3170 3171 3172 3173 3174 3175 3176 3177
  req.consumerId = tmq->consumerId;

  int32_t msgSize = tSerializeSMqSeekReq(NULL, 0, &req);
  if (msgSize < 0) {
    return TSDB_CODE_PAR_INTERNAL_ERROR;
  }

  char* msg = taosMemoryCalloc(1, msgSize);
  if (NULL == msg) {
    return TSDB_CODE_OUT_OF_MEMORY;
  }

  if (tSerializeSMqSeekReq(msg, msgSize, &req) < 0) {
    taosMemoryFree(msg);
    return TSDB_CODE_PAR_INTERNAL_ERROR;
  }

  SMsgSendInfo* sendInfo = taosMemoryCalloc(1, sizeof(SMsgSendInfo));
  if (sendInfo == NULL) {
    taosMemoryFree(msg);
    return TSDB_CODE_OUT_OF_MEMORY;
  }

  SMqSeekParam* pParam = taosMemoryMalloc(sizeof(SMqSeekParam));
  if (pParam == NULL) {
    taosMemoryFree(msg);
    taosMemoryFree(sendInfo);
    return TSDB_CODE_OUT_OF_MEMORY;
  }
  tsem_init(&pParam->sem, 0, 0);

  sendInfo->msgInfo = (SDataBuf){.pData = msg, .len = msgSize, .handle = NULL};
  sendInfo->requestId = generateRequestId();
  sendInfo->requestObjRefId = 0;
  sendInfo->param = pParam;
  sendInfo->fp = tmqSeekCb;
  sendInfo->msgType = TDMT_VND_TMQ_SEEK;
H
Haojun Liao 已提交
3178

3179
  int64_t transporterId = 0;
3180
  asyncSendMsgToServer(tmq->pTscObj->pAppInfo->pTransporter, &epSet, &transporterId, sendInfo);
3181

3182
  tsem_wait(&pParam->sem);
3183
  code = pParam->code;
3184 3185
  tsem_destroy(&pParam->sem);
  taosMemoryFree(pParam);
H
Haojun Liao 已提交
3186

3187
  tscInfo("consumer:0x%" PRIx64 "send seek to vgId:%d, return code:%s", tmq->consumerId, vgId, tstrerror(code));
3188 3189

  return code;
P
plum-lihui 已提交
3190
}