crash_gen.py 99.3 KB
Newer Older
S
Shuduo Sang 已提交
1
# -----!/usr/bin/python3.7
S
Steven Li 已提交
2 3 4 5 6 7 8 9 10 11 12 13
###################################################################
#           Copyright (c) 2016 by TAOS Technologies, Inc.
#                     All rights reserved.
#
#  This file is proprietary and confidential to TAOS Technologies.
#  No part of this file may be reproduced, stored, transmitted,
#  disclosed or used in any form or by any means other than as
#  expressly provided by the written permission from Jianhui Tao
#
###################################################################

# -*- coding: utf-8 -*-
S
Shuduo Sang 已提交
14 15 16 17 18 19 20 21 22 23 24 25 26 27
# For type hinting before definition, ref:
# https://stackoverflow.com/questions/33533148/how-do-i-specify-that-the-return-type-of-a-method-is-the-same-as-the-class-itsel
from __future__ import annotations
import taos
from util.sql import *
from util.cases import *
from util.dnodes import *
from util.log import *
from typing import Set
from typing import Dict
from typing import List
from requests.auth import HTTPBasicAuth
import textwrap
import time
28
import datetime
S
Shuduo Sang 已提交
29
import random
30
import logging
S
Shuduo Sang 已提交
31 32 33 34 35
import threading
import requests
import copy
import argparse
import getopt
36

S
Steven Li 已提交
37
import sys
38
import os
39
import signal
40
import traceback
41 42 43
import resource
from guppy import hpy
import gc
44

45 46
from .service_manager import ServiceManager, TdeInstance
from .misc import Logging, Status, CrashGenError, Dice, Helper, Progress
47

48 49 50 51
# Require Python 3
if sys.version_info[0] < 3:
    raise Exception("Must be using Python 3")

S
Shuduo Sang 已提交
52
# Global variables, tried to keep a small number.
53 54 55

# Command-line/Environment Configurations, will set a bit later
# ConfigNameSpace = argparse.Namespace
56 57
gConfig:    argparse.Namespace 
gSvcMgr:    ServiceManager # TODO: refactor this hack, use dep injection
58
# logger:     logging.Logger
59
gContainer: Container
S
Steven Li 已提交
60

61 62
# def runThread(wt: WorkerThread):
#     wt.run()
63

64

S
Steven Li 已提交
65
class WorkerThread:
66
    def __init__(self, pool: ThreadPool, tid, tc: ThreadCoordinator,
S
Shuduo Sang 已提交
67 68 69
                 # te: TaskExecutor,
                 ):  # note: main thread context!
        # self._curStep = -1
70
        self._pool = pool
S
Shuduo Sang 已提交
71 72
        self._tid = tid
        self._tc = tc  # type: ThreadCoordinator
S
Steven Li 已提交
73
        # self.threadIdent = threading.get_ident()
74 75
        # self._thread = threading.Thread(target=runThread, args=(self,))
        self._thread = threading.Thread(target=self.run)
76
        self._stepGate = threading.Event()
S
Steven Li 已提交
77

78
        # Let us have a DB connection of our own
S
Shuduo Sang 已提交
79
        if (gConfig.per_thread_db_connection):  # type: ignore
80
            # print("connector_type = {}".format(gConfig.connector_type))
81 82 83 84 85 86 87 88 89 90 91
            if gConfig.connector_type == 'native':
                self._dbConn = DbConn.createNative() 
            elif gConfig.connector_type == 'rest':
                self._dbConn = DbConn.createRest() 
            elif gConfig.connector_type == 'mixed':
                if Dice.throw(2) == 0: # 1/2 chance
                    self._dbConn = DbConn.createNative() 
                else:
                    self._dbConn = DbConn.createRest() 
            else:
                raise RuntimeError("Unexpected connector type: {}".format(gConfig.connector_type))
92

93
        # self._dbInUse = False  # if "use db" was executed already
94

95
    def logDebug(self, msg):
96
        Logging.debug("    TRD[{}] {}".format(self._tid, msg))
97 98

    def logInfo(self, msg):
99
        Logging.info("    TRD[{}] {}".format(self._tid, msg))
100

101 102
    # def dbInUse(self):
    #     return self._dbInUse
103

104 105 106 107
    # def useDb(self):
    #     if (not self._dbInUse):
    #         self.execSql("use db")
    #     self._dbInUse = True
108

109
    def getTaskExecutor(self):
S
Shuduo Sang 已提交
110
        return self._tc.getTaskExecutor()
111

S
Steven Li 已提交
112
    def start(self):
113
        self._thread.start()  # AFTER the thread is recorded
S
Steven Li 已提交
114

S
Shuduo Sang 已提交
115
    def run(self):
S
Steven Li 已提交
116
        # initialization after thread starts, in the thread context
117
        # self.isSleeping = False
118
        Logging.info("Starting to run thread: {}".format(self._tid))
119

S
Shuduo Sang 已提交
120
        if (gConfig.per_thread_db_connection):  # type: ignore
121
            Logging.debug("Worker thread openning database connection")
122
            self._dbConn.open()
S
Steven Li 已提交
123

S
Shuduo Sang 已提交
124 125
        self._doTaskLoop()

126
        # clean up
S
Shuduo Sang 已提交
127
        if (gConfig.per_thread_db_connection):  # type: ignore
128 129 130
            if self._dbConn.isOpen: #sometimes it is not open
                self._dbConn.close()
            else:
131
                Logging.warning("Cleaning up worker thread, dbConn already closed")
132

S
Shuduo Sang 已提交
133
    def _doTaskLoop(self):
134 135
        # while self._curStep < self._pool.maxSteps:
        # tc = ThreadCoordinator(None)
S
Shuduo Sang 已提交
136 137
        while True:
            tc = self._tc  # Thread Coordinator, the overall master
138 139 140
            try:
                tc.crossStepBarrier()  # shared barrier first, INCLUDING the last one
            except threading.BrokenBarrierError as err: # main thread timed out
141
                print("_bto", end="")
142
                Logging.debug("[TRD] Worker thread exiting due to main thread barrier time-out")
143 144
                break

145
            Logging.debug("[TRD] Worker thread [{}] exited barrier...".format(self._tid))
146
            self.crossStepGate()   # then per-thread gate, after being tapped
147
            Logging.debug("[TRD] Worker thread [{}] exited step gate...".format(self._tid))
148
            if not self._tc.isRunning():
149
                print("_wts", end="")
150
                Logging.debug("[TRD] Thread Coordinator not running any more, worker thread now stopping...")
151 152
                break

153
            # Before we fetch the task and run it, let's ensure we properly "use" the database (not needed any more)
154
            try:
155 156 157
                if (gConfig.per_thread_db_connection):  # most likely TRUE
                    if not self._dbConn.isOpen:  # might have been closed during server auto-restart
                        self._dbConn.open()
158
                # self.useDb() # might encounter exceptions. TODO: catch
159 160
            except taos.error.ProgrammingError as err:
                errno = Helper.convertErrno(err.errno)
161
                if errno in [0x383, 0x386, 0x00B, 0x014]  : # invalid database, dropping, Unable to establish connection, Database not ready
162 163 164 165 166 167
                    # ignore
                    dummy = 0
                else:
                    print("\nCaught programming error. errno=0x{:X}, msg={} ".format(errno, err.msg))
                    raise

168
            # Fetch a task from the Thread Coordinator
169
            Logging.debug( "[TRD] Worker thread [{}] about to fetch task".format(self._tid))
170
            task = tc.fetchTask()
171 172

            # Execute such a task
173
            Logging.debug("[TRD] Worker thread [{}] about to execute task: {}".format(
S
Shuduo Sang 已提交
174
                    self._tid, task.__class__.__name__))
175
            task.execute(self)
176
            tc.saveExecutedTask(task)
177
            Logging.debug("[TRD] Worker thread [{}] finished executing task".format(self._tid))
S
Shuduo Sang 已提交
178

179
            # self._dbInUse = False  # there may be changes between steps
180
        # print("_wtd", end=None) # worker thread died
181

S
Shuduo Sang 已提交
182 183
    def verifyThreadSelf(self):  # ensure we are called by this own thread
        if (threading.get_ident() != self._thread.ident):
S
Steven Li 已提交
184 185
            raise RuntimeError("Unexpectly called from other threads")

S
Shuduo Sang 已提交
186 187
    def verifyThreadMain(self):  # ensure we are called by the main thread
        if (threading.get_ident() != threading.main_thread().ident):
S
Steven Li 已提交
188 189 190
            raise RuntimeError("Unexpectly called from other threads")

    def verifyThreadAlive(self):
S
Shuduo Sang 已提交
191
        if (not self._thread.is_alive()):
S
Steven Li 已提交
192 193
            raise RuntimeError("Unexpected dead thread")

194
    # A gate is different from a barrier in that a thread needs to be "tapped"
S
Steven Li 已提交
195 196
    def crossStepGate(self):
        self.verifyThreadAlive()
S
Shuduo Sang 已提交
197 198
        self.verifyThreadSelf()  # only allowed by ourselves

199
        # Wait again at the "gate", waiting to be "tapped"
200
        Logging.debug(
S
Shuduo Sang 已提交
201 202 203
            "[TRD] Worker thread {} about to cross the step gate".format(
                self._tid))
        self._stepGate.wait()
204
        self._stepGate.clear()
S
Shuduo Sang 已提交
205

206
        # self._curStep += 1  # off to a new step...
S
Steven Li 已提交
207

S
Shuduo Sang 已提交
208
    def tapStepGate(self):  # give it a tap, release the thread waiting there
209
        # self.verifyThreadAlive()
S
Shuduo Sang 已提交
210 211
        self.verifyThreadMain()  # only allowed for main thread

212
        if self._thread.is_alive():
213
            Logging.debug("[TRD] Tapping worker thread {}".format(self._tid))
214 215 216 217
            self._stepGate.set()  # wake up!
            time.sleep(0)  # let the released thread run a bit
        else:
            print("_tad", end="") # Thread already dead
218

S
Shuduo Sang 已提交
219
    def execSql(self, sql):  # TODO: expose DbConn directly
220
        return self.getDbConn().execute(sql)
221

S
Shuduo Sang 已提交
222
    def querySql(self, sql):  # TODO: expose DbConn directly
223
        return self.getDbConn().query(sql)
224 225

    def getQueryResult(self):
226
        return self.getDbConn().getQueryResult()
227

228
    def getDbConn(self) -> DbConn :
S
Shuduo Sang 已提交
229 230
        if (gConfig.per_thread_db_connection):
            return self._dbConn
231
        else:
232
            return self._tc.getDbManager().getDbConn()
233

234 235
    # def querySql(self, sql): # not "execute", since we are out side the DB context
    #     if ( gConfig.per_thread_db_connection ):
S
Shuduo Sang 已提交
236
    #         return self._dbConn.query(sql)
237 238
    #     else:
    #         return self._tc.getDbState().getDbConn().query(sql)
239

240
# The coordinator of all worker threads, mostly running in main thread
S
Shuduo Sang 已提交
241 242


243
class ThreadCoordinator:
244
    WORKER_THREAD_TIMEOUT = 180 # one minute
245

246
    def __init__(self, pool: ThreadPool, dbManager: DbManager):
S
Shuduo Sang 已提交
247
        self._curStep = -1  # first step is 0
248
        self._pool = pool
249
        # self._wd = wd
S
Shuduo Sang 已提交
250
        self._te = None  # prepare for every new step
251
        self._dbManager = dbManager
S
Shuduo Sang 已提交
252 253
        self._executedTasks: List[Task] = []  # in a given step
        self._lock = threading.RLock()  # sync access for a few things
S
Steven Li 已提交
254

S
Shuduo Sang 已提交
255 256
        self._stepBarrier = threading.Barrier(
            self._pool.numThreads + 1)  # one barrier for all threads
257
        self._execStats = ExecutionStats()
258
        self._runStatus = Status.STATUS_RUNNING
259
        self._initDbs()
S
Steven Li 已提交
260

261 262 263
    def getTaskExecutor(self):
        return self._te

S
Shuduo Sang 已提交
264
    def getDbManager(self) -> DbManager:
265
        return self._dbManager
266

267 268
    def crossStepBarrier(self, timeout=None):
        self._stepBarrier.wait(timeout) 
269

270
    def requestToStop(self):
271
        self._runStatus = Status.STATUS_STOPPING
272 273
        self._execStats.registerFailure("User Interruption")

274
    def _runShouldEnd(self, transitionFailed, hasAbortedTask, workerTimeout):
275 276 277
        maxSteps = gConfig.max_steps  # type: ignore
        if self._curStep >= (maxSteps - 1): # maxStep==10, last curStep should be 9
            return True
278
        if self._runStatus != Status.STATUS_RUNNING:
279 280 281 282 283
            return True
        if transitionFailed:
            return True
        if hasAbortedTask:
            return True
284 285
        if workerTimeout:
            return True
286 287 288 289 290 291 292 293 294 295 296 297 298
        return False

    def _hasAbortedTask(self): # from execution of previous step
        for task in self._executedTasks:
            if task.isAborted():
                # print("Task aborted: {}".format(task))
                # hasAbortedTask = True
                return True
        return False

    def _releaseAllWorkerThreads(self, transitionFailed):
        self._curStep += 1  # we are about to get into next step. TODO: race condition here!
        # Now not all threads had time to go to sleep
299
        Logging.debug(
300 301 302 303 304 305 306
            "--\r\n\n--> Step {} starts with main thread waking up".format(self._curStep))

        # A new TE for the new step
        self._te = None # set to empty first, to signal worker thread to stop
        if not transitionFailed:  # only if not failed
            self._te = TaskExecutor(self._curStep)

307
        Logging.debug("[TRD] Main thread waking up at step {}, tapping worker threads".format(
308 309
                self._curStep))  # Now not all threads had time to go to sleep
        # Worker threads will wake up at this point, and each execute it's own task
310
        self.tapAllThreads() # release all worker thread from their "gates"
311 312 313 314 315

    def _syncAtBarrier(self):
         # Now main thread (that's us) is ready to enter a step
        # let other threads go past the pool barrier, but wait at the
        # thread gate
316
        Logging.debug("[TRD] Main thread about to cross the barrier")
317
        self.crossStepBarrier(timeout=self.WORKER_THREAD_TIMEOUT)
318
        self._stepBarrier.reset()  # Other worker threads should now be at the "gate"
319
        Logging.debug("[TRD] Main thread finished crossing the barrier")
320 321 322 323

    def _doTransition(self):
        transitionFailed = False
        try:
324 325 326
            for x in self._dbs:
                db = x # type: Database
                sm = db.getStateMachine()
327
                Logging.debug("[STT] starting transitions for DB: {}".format(db.getName()))
328 329 330
                # at end of step, transiton the DB state
                tasksForDb = db.filterTasks(self._executedTasks)
                sm.transition(tasksForDb, self.getDbManager().getDbConn())
331
                Logging.debug("[STT] transition ended for DB: {}".format(db.getName()))
332 333

            # Due to limitation (or maybe not) of the TD Python library,
334
            # we cannot share connections across threads
335 336 337 338
            # Here we are in main thread, we cannot operate the connections created in workers
            # Moving below to task loop
            # if sm.hasDatabase():
            #     for t in self._pool.threadList:
339
            #         Logging.debug("[DB] use db for all worker threads")
340
            #         t.useDb()
341 342
                    # t.execSql("use db") # main thread executing "use
                    # db" on behalf of every worker thread
343

344 345
        except taos.error.ProgrammingError as err:
            if (err.msg == 'network unavailable'):  # broken DB connection
346
                Logging.info("DB connection broken, execution failed")
347 348 349 350 351 352 353 354
                traceback.print_stack()
                transitionFailed = True
                self._te = None  # Not running any more
                self._execStats.registerFailure("Broken DB Connection")
                # continue # don't do that, need to tap all threads at
                # end, and maybe signal them to stop
            else:
                raise
S
Steven Li 已提交
355
        # return transitionFailed # Why did we have this??!!
356 357 358

        self.resetExecutedTasks()  # clear the tasks after we are done
        # Get ready for next step
359
        Logging.debug("<-- Step {} finished, trasition failed = {}".format(self._curStep, transitionFailed))
360 361
        return transitionFailed

S
Shuduo Sang 已提交
362
    def run(self):
363
        self._pool.createAndStartThreads(self)
S
Steven Li 已提交
364 365

        # Coordinate all threads step by step
S
Shuduo Sang 已提交
366
        self._curStep = -1  # not started yet
367
        
S
Shuduo Sang 已提交
368
        self._execStats.startExec()  # start the stop watch
369 370
        transitionFailed = False
        hasAbortedTask = False
371 372
        workerTimeout = False
        while not self._runShouldEnd(transitionFailed, hasAbortedTask, workerTimeout):
373 374 375
            if not gConfig.debug: # print this only if we are not in debug mode    
                Progress.emit(Progress.STEP_BOUNDARY)            
                # print(".", end="", flush=True)
376 377 378 379 380 381 382 383
            # if (self._curStep % 2) == 0: # print memory usage once every 10 steps
            #     memUsage = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
            #     print("[m:{}]".format(memUsage), end="", flush=True) # print memory usage
            # if (self._curStep % 10) == 3: 
            #     h = hpy()
            #     print("\n")        
            #     print(h.heap())
            
384
                        
385 386
            try:
                self._syncAtBarrier() # For now just cross the barrier
387
                Progress.emit(Progress.END_THREAD_STEP)
388
            except threading.BrokenBarrierError as err:
389
                Logging.info("Main loop aborted, caused by worker thread time-out")
390 391 392 393 394 395
                self._execStats.registerFailure("Aborted due to worker thread timeout")
                print("\n\nWorker Thread time-out detected, important thread info:")
                ts = ThreadStacks()
                ts.print(filterInternal=True)
                workerTimeout = True
                break
396 397

            # At this point, all threads should be pass the overall "barrier" and before the per-thread "gate"
S
Shuduo Sang 已提交
398 399
            # We use this period to do house keeping work, when all worker
            # threads are QUIET.
400 401
            hasAbortedTask = self._hasAbortedTask() # from previous step
            if hasAbortedTask: 
402
                Logging.info("Aborted task encountered, exiting test program")
403
                self._execStats.registerFailure("Aborted Task Encountered")
404
                break # do transition only if tasks are error free
S
Shuduo Sang 已提交
405

406
            # Ending previous step
407 408 409 410
            try:
                transitionFailed = self._doTransition() # To start, we end step -1 first
            except taos.error.ProgrammingError as err:
                transitionFailed = True
411
                errno2 = Helper.convertErrno(err.errno)  # correct error scheme
S
Steven Li 已提交
412
                errMsg = "Transition failed: errno=0x{:X}, msg: {}".format(errno2, err)
413
                Logging.info(errMsg)
414
                traceback.print_exc()
S
Steven Li 已提交
415
                self._execStats.registerFailure(errMsg)
416

417
            # Then we move on to the next step
418
            Progress.emit(Progress.BEGIN_THREAD_STEP)
419
            self._releaseAllWorkerThreads(transitionFailed)                    
420

421
        if hasAbortedTask or transitionFailed : # abnormal ending, workers waiting at "gate"
422
            Logging.debug("Abnormal ending of main thraed")
423
        elif workerTimeout:
424
            Logging.debug("Abnormal ending of main thread, due to worker timeout")
425
        else: # regular ending, workers waiting at "barrier"
426
            Logging.debug("Regular ending, main thread waiting for all worker threads to stop...")
427
            self._syncAtBarrier()
428

429
        self._te = None  # No more executor, time to end
430
        Logging.debug("Main thread tapping all threads one last time...")
431
        self.tapAllThreads()  # Let the threads run one last time
432

433 434
        Logging.debug("\r\n\n--> Main thread ready to finish up...")
        Logging.debug("Main thread joining all threads")
S
Shuduo Sang 已提交
435
        self._pool.joinAll()  # Get all threads to finish
436
        Logging.info("\nAll worker threads finished")
437 438
        self._execStats.endExec()

439 440 441 442 443 444 445 446 447 448 449 450 451
    def cleanup(self): # free resources
        self._pool.cleanup()

        self._pool = None
        self._te = None  
        self._dbManager = None
        self._executedTasks = None
        self._lock = None
        self._stepBarrier = None
        self._execStats = None
        self._runStatus = None


452 453
    def printStats(self):
        self._execStats.printStats()
S
Steven Li 已提交
454

S
Steven Li 已提交
455 456 457 458 459 460
    def isFailed(self):
        return self._execStats.isFailed()

    def getExecStats(self):
        return self._execStats

S
Shuduo Sang 已提交
461
    def tapAllThreads(self):  # in a deterministic manner
S
Steven Li 已提交
462
        wakeSeq = []
S
Shuduo Sang 已提交
463 464
        for i in range(self._pool.numThreads):  # generate a random sequence
            if Dice.throw(2) == 1:
S
Steven Li 已提交
465 466 467
                wakeSeq.append(i)
            else:
                wakeSeq.insert(0, i)
468
        Logging.debug(
S
Shuduo Sang 已提交
469 470
            "[TRD] Main thread waking up worker threads: {}".format(
                str(wakeSeq)))
471
        # TODO: set dice seed to a deterministic value
S
Steven Li 已提交
472
        for i in wakeSeq:
S
Shuduo Sang 已提交
473 474 475
            # TODO: maybe a bit too deep?!
            self._pool.threadList[i].tapStepGate()
            time.sleep(0)  # yield
S
Steven Li 已提交
476

477
    def isRunning(self):
S
Shuduo Sang 已提交
478
        return self._te is not None
479

480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496
    def _initDbs(self):
        ''' Initialize multiple databases, invoked at __ini__() time '''
        self._dbs = [] # type: List[Database]
        dbc = self.getDbManager().getDbConn()
        if gConfig.max_dbs == 0:
            self._dbs.append(Database(0, dbc))
        else:
            for i in range(gConfig.max_dbs):
                self._dbs.append(Database(i, dbc))

    def pickDatabase(self):
        idxDb = 0
        if gConfig.max_dbs != 0 :
            idxDb = Dice.throw(gConfig.max_dbs) # 0 to N-1
        db = self._dbs[idxDb] # type: Database
        return db

S
Shuduo Sang 已提交
497
    def fetchTask(self) -> Task:
498 499 500
        ''' The thread coordinator (that's us) is responsible for fetching a task
            to be executed next.
        '''
S
Shuduo Sang 已提交
501
        if (not self.isRunning()):  # no task
502
            raise RuntimeError("Cannot fetch task when not running")
503

S
Shuduo Sang 已提交
504
        # pick a task type for current state
505 506 507
        db = self.pickDatabase()
        taskType = db.getStateMachine().pickTaskType() # type: Task
        return taskType(self._execStats, db)  # create a task from it
508 509

    def resetExecutedTasks(self):
S
Shuduo Sang 已提交
510
        self._executedTasks = []  # should be under single thread
511 512 513 514

    def saveExecutedTask(self, task):
        with self._lock:
            self._executedTasks.append(task)
515

516
class ThreadPool:
517
    def __init__(self, numThreads, maxSteps):
518 519 520 521
        self.numThreads = numThreads
        self.maxSteps = maxSteps
        # Internal class variables
        self.curStep = 0
S
Shuduo Sang 已提交
522 523
        self.threadList = []  # type: List[WorkerThread]

524
    # starting to run all the threads, in locking steps
525
    def createAndStartThreads(self, tc: ThreadCoordinator):
S
Shuduo Sang 已提交
526 527
        for tid in range(0, self.numThreads):  # Create the threads
            workerThread = WorkerThread(self, tid, tc)
528
            self.threadList.append(workerThread)
S
Shuduo Sang 已提交
529
            workerThread.start()  # start, but should block immediately before step 0
530 531 532

    def joinAll(self):
        for workerThread in self.threadList:
533
            Logging.debug("Joining thread...")
534 535
            workerThread._thread.join()

536 537 538
    def cleanup(self):
        self.threadList = None # maybe clean up each?

539 540
# A queue of continguous POSITIVE integers, used by DbManager to generate continuous numbers
# for new table names
S
Shuduo Sang 已提交
541 542


S
Steven Li 已提交
543 544
class LinearQueue():
    def __init__(self):
545
        self.firstIndex = 1  # 1st ever element
S
Steven Li 已提交
546
        self.lastIndex = 0
S
Shuduo Sang 已提交
547 548
        self._lock = threading.RLock()  # our functions may call each other
        self.inUse = set()  # the indexes that are in use right now
S
Steven Li 已提交
549

550
    def toText(self):
S
Shuduo Sang 已提交
551 552
        return "[{}..{}], in use: {}".format(
            self.firstIndex, self.lastIndex, self.inUse)
553 554

    # Push (add new element, largest) to the tail, and mark it in use
S
Shuduo Sang 已提交
555
    def push(self):
556
        with self._lock:
S
Shuduo Sang 已提交
557 558
            # if ( self.isEmpty() ):
            #     self.lastIndex = self.firstIndex
559
            #     return self.firstIndex
560 561
            # Otherwise we have something
            self.lastIndex += 1
562 563
            self.allocate(self.lastIndex)
            # self.inUse.add(self.lastIndex) # mark it in use immediately
564
            return self.lastIndex
S
Steven Li 已提交
565 566

    def pop(self):
567
        with self._lock:
S
Shuduo Sang 已提交
568 569 570 571
            if (self.isEmpty()):
                # raise RuntimeError("Cannot pop an empty queue")
                return False  # TODO: None?

572
            index = self.firstIndex
S
Shuduo Sang 已提交
573
            if (index in self.inUse):
574 575
                return False

576 577 578 579 580 581 582
            self.firstIndex += 1
            return index

    def isEmpty(self):
        return self.firstIndex > self.lastIndex

    def popIfNotEmpty(self):
583
        with self._lock:
584 585 586 587
            if (self.isEmpty()):
                return 0
            return self.pop()

S
Steven Li 已提交
588
    def allocate(self, i):
589
        with self._lock:
590
            # Logging.debug("LQ allocating item {}".format(i))
S
Shuduo Sang 已提交
591 592 593
            if (i in self.inUse):
                raise RuntimeError(
                    "Cannot re-use same index in queue: {}".format(i))
594 595
            self.inUse.add(i)

S
Steven Li 已提交
596
    def release(self, i):
597
        with self._lock:
598
            # Logging.debug("LQ releasing item {}".format(i))
S
Shuduo Sang 已提交
599
            self.inUse.remove(i)  # KeyError possible, TODO: why?
600 601 602 603

    def size(self):
        return self.lastIndex + 1 - self.firstIndex

S
Steven Li 已提交
604
    def pickAndAllocate(self):
S
Shuduo Sang 已提交
605
        if (self.isEmpty()):
606 607
            return None
        with self._lock:
S
Shuduo Sang 已提交
608
            cnt = 0  # counting the interations
609 610
            while True:
                cnt += 1
S
Shuduo Sang 已提交
611
                if (cnt > self.size() * 10):  # 10x iteration already
612 613
                    # raise RuntimeError("Failed to allocate LinearQueue element")
                    return None
S
Shuduo Sang 已提交
614 615
                ret = Dice.throwRange(self.firstIndex, self.lastIndex + 1)
                if (ret not in self.inUse):
616 617 618
                    self.allocate(ret)
                    return ret

S
Shuduo Sang 已提交
619

620
class DbConn:
621
    TYPE_NATIVE = "native-c"
622
    TYPE_REST =   "rest-api"
623 624 625 626 627 628 629 630 631
    TYPE_INVALID = "invalid"

    @classmethod
    def create(cls, connType):
        if connType == cls.TYPE_NATIVE:
            return DbConnNative()
        elif connType == cls.TYPE_REST:
            return DbConnRest()
        else:
S
Shuduo Sang 已提交
632 633
            raise RuntimeError(
                "Unexpected connection type: {}".format(connType))
634 635 636 637 638 639 640 641 642

    @classmethod
    def createNative(cls):
        return cls.create(cls.TYPE_NATIVE)

    @classmethod
    def createRest(cls):
        return cls.create(cls.TYPE_REST)

643 644
    def __init__(self):
        self.isOpen = False
645
        self._type = self.TYPE_INVALID
646 647 648 649
        self._lastSql = None

    def getLastSql(self):
        return self._lastSql
650 651

    def open(self):
S
Shuduo Sang 已提交
652
        if (self.isOpen):
653 654
            raise RuntimeError("Cannot re-open an existing DB connection")

655 656
        # below implemented by child classes
        self.openByType()
657

658
        Logging.debug("[DB] data connection opened, type = {}".format(self._type))
659 660
        self.isOpen = True

661 662 663
    def close(self):
        raise RuntimeError("Unexpected execution, should be overriden")

S
Shuduo Sang 已提交
664
    def queryScalar(self, sql) -> int:
665 666
        return self._queryAny(sql)

S
Shuduo Sang 已提交
667
    def queryString(self, sql) -> str:
668 669
        return self._queryAny(sql)

S
Shuduo Sang 已提交
670 671
    def _queryAny(self, sql):  # actual query result as an int
        if (not self.isOpen):
672
            raise RuntimeError("Cannot query database until connection is open")
673
        nRows = self.query(sql)
S
Shuduo Sang 已提交
674
        if nRows != 1:
675 676 677 678
            raise taos.error.ProgrammingError(
                "Unexpected result for query: {}, rows = {}".format(sql, nRows), 
                (0x991 if nRows==0 else 0x992)
            )
679
        if self.getResultRows() != 1 or self.getResultCols() != 1:
680
            raise RuntimeError("Unexpected result set for query: {}".format(sql))
681 682
        return self.getQueryResult()[0][0]

683 684 685
    def use(self, dbName):
        self.execute("use {}".format(dbName))

686 687 688 689 690 691 692
    def existsDatabase(self, dbName: str):
        ''' Check if a certain database exists '''
        self.query("show databases")
        dbs = [v[0] for v in self.getQueryResult()] # ref: https://stackoverflow.com/questions/643823/python-list-transformation
        # ret2 = dbName in dbs
        # print("dbs = {}, str = {}, ret2={}, type2={}".format(dbs, dbName,ret2, type(dbName)))
        return dbName in dbs # TODO: super weird type mangling seen, once here
693 694 695 696

    def hasTables(self):
        return self.query("show tables") > 0

697
    def execute(self, sql):
698
        ''' Return the number of rows affected'''
699
        raise RuntimeError("Unexpected execution, should be overriden")
S
Shuduo Sang 已提交
700

701 702 703 704 705 706 707 708 709
    def safeExecute(self, sql):
        '''Safely execute any SQL query, returning True/False upon success/failure'''
        try:
            self.execute(sql)
            return True # ignore num of results, return success
        except taos.error.ProgrammingError as err:
            return False # failed, for whatever TAOS reason
        # Not possile to reach here, non-TAOS exception would have been thrown

710
    def query(self, sql) -> int: # return num rows returned
711
        ''' Return the number of rows affected'''
712 713
        raise RuntimeError("Unexpected execution, should be overriden")

714 715
    def openByType(self):
        raise RuntimeError("Unexpected execution, should be overriden")
S
Shuduo Sang 已提交
716

717 718
    def getQueryResult(self):
        raise RuntimeError("Unexpected execution, should be overriden")
S
Shuduo Sang 已提交
719

720 721
    def getResultRows(self):
        raise RuntimeError("Unexpected execution, should be overriden")
S
Shuduo Sang 已提交
722

723 724 725 726
    def getResultCols(self):
        raise RuntimeError("Unexpected execution, should be overriden")

# Sample: curl -u root:taosdata -d "show databases" localhost:6020/rest/sql
S
Shuduo Sang 已提交
727 728


729 730 731 732
class DbConnRest(DbConn):
    def __init__(self):
        super().__init__()
        self._type = self.TYPE_REST
S
Steven Li 已提交
733
        self._url = "http://localhost:6041/rest/sql"  # fixed for now
734 735
        self._result = None

S
Shuduo Sang 已提交
736 737 738
    def openByType(self):  # Open connection
        pass  # do nothing, always open

739
    def close(self):
S
Shuduo Sang 已提交
740
        if (not self.isOpen):
741
            raise RuntimeError("Cannot clean up database until connection is open")
742
        # Do nothing for REST
743
        Logging.debug("[DB] REST Database connection closed")
744 745 746
        self.isOpen = False

    def _doSql(self, sql):
747
        self._lastSql = sql # remember this, last SQL attempted
748 749 750
        try:
            r = requests.post(self._url, 
                data = sql,
751
                auth = HTTPBasicAuth('root', 'taosdata'))         
752 753 754
        except:
            print("REST API Failure (TODO: more info here)")
            raise
755 756
        rj = r.json()
        # Sanity check for the "Json Result"
S
Shuduo Sang 已提交
757
        if ('status' not in rj):
758 759
            raise RuntimeError("No status in REST response")

S
Shuduo Sang 已提交
760 761 762 763
        if rj['status'] == 'error':  # clearly reported error
            if ('code' not in rj):  # error without code
                raise RuntimeError("REST error return without code")
            errno = rj['code']  # May need to massage this in the future
764
            # print("Raising programming error with REST return: {}".format(rj))
S
Shuduo Sang 已提交
765 766
            raise taos.error.ProgrammingError(
                rj['desc'], errno)  # todo: check existance of 'desc'
767

S
Shuduo Sang 已提交
768 769 770 771
        if rj['status'] != 'succ':  # better be this
            raise RuntimeError(
                "Unexpected REST return status: {}".format(
                    rj['status']))
772 773

        nRows = rj['rows'] if ('rows' in rj) else 0
S
Shuduo Sang 已提交
774
        self._result = rj
775 776
        return nRows

S
Shuduo Sang 已提交
777 778 779 780
    def execute(self, sql):
        if (not self.isOpen):
            raise RuntimeError(
                "Cannot execute database commands until connection is open")
781
        Logging.debug("[SQL-REST] Executing SQL: {}".format(sql))
782
        nRows = self._doSql(sql)
783
        Logging.debug(
S
Shuduo Sang 已提交
784
            "[SQL-REST] Execution Result, nRows = {}, SQL = {}".format(nRows, sql))
785 786
        return nRows

S
Shuduo Sang 已提交
787
    def query(self, sql):  # return rows affected
788 789 790 791 792 793 794 795 796 797 798 799 800
        return self.execute(sql)

    def getQueryResult(self):
        return self._result['data']

    def getResultRows(self):
        print(self._result)
        raise RuntimeError("TBD")
        # return self._tdSql.queryRows

    def getResultCols(self):
        print(self._result)
        raise RuntimeError("TBD")
S
Shuduo Sang 已提交
801

802
    # Duplicate code from TDMySQL, TODO: merge all this into DbConnNative
803 804


805
class MyTDSql:
806 807 808 809 810 811 812
    # Class variables
    _clsLock = threading.Lock() # class wide locking
    longestQuery = None # type: str
    longestQueryTime = 0.0 # seconds
    lqStartTime = 0.0
    # lqEndTime = 0.0 # Not needed, as we have the two above already

813 814 815 816 817
    def __init__(self, hostAddr, cfgPath):
        # Make the DB connection
        self._conn = taos.connect(host=hostAddr, config=cfgPath) 
        self._cursor = self._conn.cursor()

818 819 820 821
        self.queryRows = 0
        self.queryCols = 0
        self.affectedRows = 0

822 823
    # def init(self, cursor, log=True):
    #     self.cursor = cursor
824 825 826 827 828
        # if (log):
        #     caller = inspect.getframeinfo(inspect.stack()[1][0])
        #     self.cursor.log(caller.filename + ".sql")

    def close(self):
829
        self._cursor.close() # can we double close?
830 831
        self._conn.close() # TODO: very important, cursor close does NOT close DB connection!
        self._cursor.close()
832

833 834 835
    def _execInternal(self, sql):
        startTime = time.time() 
        ret = self._cursor.execute(sql)
836
        # print("\nSQL success: {}".format(sql))
837 838 839 840 841 842 843 844 845 846
        queryTime =  time.time() - startTime
        # Record the query time
        cls = self.__class__
        if queryTime > (cls.longestQueryTime + 0.01) :
            with cls._clsLock:
                cls.longestQuery = sql
                cls.longestQueryTime = queryTime
                cls.lqStartTime = startTime
        return ret

847 848 849
    def query(self, sql):
        self.sql = sql
        try:
850
            self._execInternal(sql)
851
            self.queryResult = self._cursor.fetchall()
852
            self.queryRows = len(self.queryResult)
853
            self.queryCols = len(self._cursor.description)
854 855 856 857 858 859
        except Exception as e:
            # caller = inspect.getframeinfo(inspect.stack()[1][0])
            # args = (caller.filename, caller.lineno, sql, repr(e))
            # tdLog.exit("%s(%d) failed: sql:%s, %s" % args)
            raise
        return self.queryRows
860

861 862 863
    def execute(self, sql):
        self.sql = sql
        try:
864
            self.affectedRows = self._execInternal(sql)
865 866 867 868 869 870 871
        except Exception as e:
            # caller = inspect.getframeinfo(inspect.stack()[1][0])
            # args = (caller.filename, caller.lineno, sql, repr(e))
            # tdLog.exit("%s(%d) failed: sql:%s, %s" % args)
            raise
        return self.affectedRows

872 873 874 875 876 877 878 879 880 881 882 883
class DbConnNative(DbConn):
    # Class variables
    _lock = threading.Lock()
    _connInfoDisplayed = False
    totalConnections = 0 # Not private

    def __init__(self):
        super().__init__()
        self._type = self.TYPE_NATIVE
        self._conn = None
        # self._cursor = None        

S
Shuduo Sang 已提交
884
    def openByType(self):  # Open connection
885 886 887 888 889
        global gContainer
        tdeInstance = gContainer.defTdeInstance # set up in ClientManager, type: TdeInstance
        # cfgPath = self.getBuildPath() + "/test/cfg"
        cfgPath  = tdeInstance.getCfgDir()
        hostAddr = tdeInstance.getHostAddr()
890

891 892 893 894
        cls = self.__class__ # Get the class, to access class variables
        with cls._lock: # force single threading for opening DB connections. # TODO: whaaat??!!!
            if not cls._connInfoDisplayed:
                cls._connInfoDisplayed = True # updating CLASS variable
895
                Logging.info("Initiating TAOS native connection to {}, using config at {}".format(hostAddr, cfgPath))                    
896 897 898 899 900 901
            # Make the connection         
            # self._conn = taos.connect(host=hostAddr, config=cfgPath)  # TODO: make configurable
            # self._cursor = self._conn.cursor()
            # Record the count in the class
            self._tdSql = MyTDSql(hostAddr, cfgPath) # making DB connection
            cls.totalConnections += 1 
902
        
903
        self._tdSql.execute('reset query cache')
S
Shuduo Sang 已提交
904
        # self._cursor.execute('use db') # do this at the beginning of every
905 906

        # Open connection
907 908 909
        # self._tdSql = MyTDSql()
        # self._tdSql.init(self._cursor)
        
910
    def close(self):
S
Shuduo Sang 已提交
911
        if (not self.isOpen):
912
            raise RuntimeError("Cannot clean up database until connection is open")
913
        self._tdSql.close()
914 915 916 917 918
        # Decrement the class wide counter
        cls = self.__class__ # Get the class, to access class variables
        with cls._lock:
            cls.totalConnections -= 1

919
        Logging.debug("[DB] Database connection closed")
920
        self.isOpen = False
S
Steven Li 已提交
921

S
Shuduo Sang 已提交
922 923
    def execute(self, sql):
        if (not self.isOpen):
924
            raise RuntimeError("Cannot execute database commands until connection is open")
925
        Logging.debug("[SQL] Executing SQL: {}".format(sql))
926
        self._lastSql = sql
927
        nRows = self._tdSql.execute(sql)
928
        Logging.debug(
S
Shuduo Sang 已提交
929 930
            "[SQL] Execution Result, nRows = {}, SQL = {}".format(
                nRows, sql))
931
        return nRows
S
Steven Li 已提交
932

S
Shuduo Sang 已提交
933 934 935 936
    def query(self, sql):  # return rows affected
        if (not self.isOpen):
            raise RuntimeError(
                "Cannot query database until connection is open")
937
        Logging.debug("[SQL] Executing SQL: {}".format(sql))
938
        self._lastSql = sql
939
        nRows = self._tdSql.query(sql)
940
        Logging.debug(
S
Shuduo Sang 已提交
941 942
            "[SQL] Query Result, nRows = {}, SQL = {}".format(
                nRows, sql))
943
        return nRows
944
        # results are in: return self._tdSql.queryResult
945

946 947 948
    def getQueryResult(self):
        return self._tdSql.queryResult

949 950
    def getResultRows(self):
        return self._tdSql.queryRows
951

952 953
    def getResultCols(self):
        return self._tdSql.queryCols
954

S
Shuduo Sang 已提交
955

956
class AnyState:
S
Shuduo Sang 已提交
957 958 959
    STATE_INVALID = -1
    STATE_EMPTY = 0  # nothing there, no even a DB
    STATE_DB_ONLY = 1  # we have a DB, but nothing else
960
    STATE_TABLE_ONLY = 2  # we have a table, but totally empty
S
Shuduo Sang 已提交
961
    STATE_HAS_DATA = 3  # we have some data in the table
962 963 964 965
    _stateNames = ["Invalid", "Empty", "DB_Only", "Table_Only", "Has_Data"]

    STATE_VAL_IDX = 0
    CAN_CREATE_DB = 1
966 967 968
    # For below, if we can "drop the DB", but strictly speaking 
    # only "under normal circumstances", as we may override it with the -b option
    CAN_DROP_DB = 2  
969 970
    CAN_CREATE_FIXED_SUPER_TABLE = 3
    CAN_DROP_FIXED_SUPER_TABLE = 4
971 972 973 974 975 976 977
    CAN_ADD_DATA = 5
    CAN_READ_DATA = 6

    def __init__(self):
        self._info = self.getInfo()

    def __str__(self):
S
Shuduo Sang 已提交
978 979
        # -1 hack to accomodate the STATE_INVALID case
        return self._stateNames[self._info[self.STATE_VAL_IDX] + 1]
980

981 982
    # Each sub state tells us the "info", about itself, so we can determine
    # on things like canDropDB()
983 984 985
    def getInfo(self):
        raise RuntimeError("Must be overriden by child classes")

S
Steven Li 已提交
986 987 988 989 990 991
    def equals(self, other):
        if isinstance(other, int):
            return self.getValIndex() == other
        elif isinstance(other, AnyState):
            return self.getValIndex() == other.getValIndex()
        else:
S
Shuduo Sang 已提交
992 993 994
            raise RuntimeError(
                "Unexpected comparison, type = {}".format(
                    type(other)))
S
Steven Li 已提交
995

996 997 998
    def verifyTasksToState(self, tasks, newState):
        raise RuntimeError("Must be overriden by child classes")

S
Steven Li 已提交
999 1000 1001
    def getValIndex(self):
        return self._info[self.STATE_VAL_IDX]

1002 1003
    def getValue(self):
        return self._info[self.STATE_VAL_IDX]
S
Shuduo Sang 已提交
1004

1005 1006
    def canCreateDb(self):
        return self._info[self.CAN_CREATE_DB]
S
Shuduo Sang 已提交
1007

1008
    def canDropDb(self):
1009 1010 1011 1012
        # If user requests to run up to a number of DBs,
        # we'd then not do drop_db operations any more
        if gConfig.max_dbs > 0 : 
            return False
1013
        return self._info[self.CAN_DROP_DB]
S
Shuduo Sang 已提交
1014

1015 1016
    def canCreateFixedSuperTable(self):
        return self._info[self.CAN_CREATE_FIXED_SUPER_TABLE]
S
Shuduo Sang 已提交
1017

1018 1019
    def canDropFixedSuperTable(self):
        return self._info[self.CAN_DROP_FIXED_SUPER_TABLE]
S
Shuduo Sang 已提交
1020

1021 1022
    def canAddData(self):
        return self._info[self.CAN_ADD_DATA]
S
Shuduo Sang 已提交
1023

1024 1025 1026 1027 1028
    def canReadData(self):
        return self._info[self.CAN_READ_DATA]

    def assertAtMostOneSuccess(self, tasks, cls):
        sCnt = 0
S
Shuduo Sang 已提交
1029
        for task in tasks:
1030 1031 1032
            if not isinstance(task, cls):
                continue
            if task.isSuccess():
S
Steven Li 已提交
1033
                # task.logDebug("Task success found")
1034
                sCnt += 1
S
Shuduo Sang 已提交
1035 1036 1037
                if (sCnt >= 2):
                    raise RuntimeError(
                        "Unexpected more than 1 success with task: {}".format(cls))
1038 1039 1040 1041

    def assertIfExistThenSuccess(self, tasks, cls):
        sCnt = 0
        exists = False
S
Shuduo Sang 已提交
1042
        for task in tasks:
1043 1044
            if not isinstance(task, cls):
                continue
S
Shuduo Sang 已提交
1045
            exists = True  # we have a valid instance
1046 1047
            if task.isSuccess():
                sCnt += 1
S
Shuduo Sang 已提交
1048
        if (exists and sCnt <= 0):
S
Steven Li 已提交
1049 1050
            raise RuntimeError("Unexpected zero success for task type: {}, from tasks: {}"
                .format(cls, tasks))
1051 1052

    def assertNoTask(self, tasks, cls):
S
Shuduo Sang 已提交
1053
        for task in tasks:
1054
            if isinstance(task, cls):
S
Shuduo Sang 已提交
1055 1056
                raise CrashGenError(
                    "This task: {}, is not expected to be present, given the success/failure of others".format(cls.__name__))
1057 1058

    def assertNoSuccess(self, tasks, cls):
S
Shuduo Sang 已提交
1059
        for task in tasks:
1060 1061
            if isinstance(task, cls):
                if task.isSuccess():
S
Shuduo Sang 已提交
1062 1063
                    raise RuntimeError(
                        "Unexpected successful task: {}".format(cls))
1064 1065

    def hasSuccess(self, tasks, cls):
S
Shuduo Sang 已提交
1066
        for task in tasks:
1067 1068 1069 1070 1071 1072
            if not isinstance(task, cls):
                continue
            if task.isSuccess():
                return True
        return False

S
Steven Li 已提交
1073
    def hasTask(self, tasks, cls):
S
Shuduo Sang 已提交
1074
        for task in tasks:
S
Steven Li 已提交
1075 1076 1077 1078
            if isinstance(task, cls):
                return True
        return False

S
Shuduo Sang 已提交
1079

1080 1081 1082 1083
class StateInvalid(AnyState):
    def getInfo(self):
        return [
            self.STATE_INVALID,
S
Shuduo Sang 已提交
1084 1085 1086
            False, False,  # can create/drop Db
            False, False,  # can create/drop fixed table
            False, False,  # can insert/read data with fixed table
1087 1088 1089 1090
        ]

    # def verifyTasksToState(self, tasks, newState):

S
Shuduo Sang 已提交
1091

1092 1093 1094 1095
class StateEmpty(AnyState):
    def getInfo(self):
        return [
            self.STATE_EMPTY,
S
Shuduo Sang 已提交
1096 1097 1098
            True, False,  # can create/drop Db
            False, False,  # can create/drop fixed table
            False, False,  # can insert/read data with fixed table
1099 1100
        ]

S
Shuduo Sang 已提交
1101 1102
    def verifyTasksToState(self, tasks, newState):
        if (self.hasSuccess(tasks, TaskCreateDb)
1103
                ):  # at EMPTY, if there's succes in creating DB
S
Shuduo Sang 已提交
1104 1105 1106 1107
            if (not self.hasTask(tasks, TaskDropDb)):  # and no drop_db tasks
                # we must have at most one. TODO: compare numbers
                self.assertAtMostOneSuccess(tasks, TaskCreateDb)

1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118

class StateDbOnly(AnyState):
    def getInfo(self):
        return [
            self.STATE_DB_ONLY,
            False, True,
            True, False,
            False, False,
        ]

    def verifyTasksToState(self, tasks, newState):
S
Shuduo Sang 已提交
1119 1120 1121
        if (not self.hasTask(tasks, TaskCreateDb)):
            # only if we don't create any more
            self.assertAtMostOneSuccess(tasks, TaskDropDb)
1122 1123 1124 1125 1126

        # TODO: restore the below, the problem exists, although unlikely in real-world
        # if (gSvcMgr!=None) and gSvcMgr.isRestarting():     
        # if (gSvcMgr == None) or (not gSvcMgr.isRestarting()) : 
        #     self.assertIfExistThenSuccess(tasks, TaskDropDb)       
1127

S
Shuduo Sang 已提交
1128

1129
class StateSuperTableOnly(AnyState):
1130 1131 1132 1133 1134 1135 1136 1137 1138
    def getInfo(self):
        return [
            self.STATE_TABLE_ONLY,
            False, True,
            False, True,
            True, True,
        ]

    def verifyTasksToState(self, tasks, newState):
S
Shuduo Sang 已提交
1139
        if (self.hasSuccess(tasks, TaskDropSuperTable)
1140
                ):  # we are able to drop the table
1141
            #self.assertAtMostOneSuccess(tasks, TaskDropSuperTable)
S
Shuduo Sang 已提交
1142 1143
            # we must have had recreted it
            self.hasSuccess(tasks, TaskCreateSuperTable)
1144

1145
            # self._state = self.STATE_DB_ONLY
S
Steven Li 已提交
1146 1147
        # elif ( self.hasSuccess(tasks, AddFixedDataTask) ): # no success dropping the table, but added data
        #     self.assertNoTask(tasks, DropFixedTableTask) # not true in massively parrallel cases
1148
            # self._state = self.STATE_HAS_DATA
S
Steven Li 已提交
1149 1150 1151
        # elif ( self.hasSuccess(tasks, ReadFixedDataTask) ): # no success in prev cases, but was able to read data
            # self.assertNoTask(tasks, DropFixedTableTask)
            # self.assertNoTask(tasks, AddFixedDataTask)
1152
            # self._state = self.STATE_TABLE_ONLY # no change
S
Steven Li 已提交
1153 1154 1155
        # else: # did not drop table, did not insert data, did not read successfully, that is impossible
        #     raise RuntimeError("Unexpected no-success scenarios")
        # TODO: need to revamp!!
1156

S
Shuduo Sang 已提交
1157

1158 1159 1160 1161 1162 1163 1164 1165 1166 1167
class StateHasData(AnyState):
    def getInfo(self):
        return [
            self.STATE_HAS_DATA,
            False, True,
            False, True,
            True, True,
        ]

    def verifyTasksToState(self, tasks, newState):
S
Shuduo Sang 已提交
1168
        if (newState.equals(AnyState.STATE_EMPTY)):
1169
            self.hasSuccess(tasks, TaskDropDb)
S
Shuduo Sang 已提交
1170 1171 1172 1173
            if (not self.hasTask(tasks, TaskCreateDb)):
                self.assertAtMostOneSuccess(tasks, TaskDropDb)  # TODO: dicy
        elif (newState.equals(AnyState.STATE_DB_ONLY)):  # in DB only
            if (not self.hasTask(tasks, TaskCreateDb)
1174
                ):  # without a create_db task
S
Shuduo Sang 已提交
1175 1176
                # we must have drop_db task
                self.assertNoTask(tasks, TaskDropDb)
1177
            self.hasSuccess(tasks, TaskDropSuperTable)
1178
            # self.assertAtMostOneSuccess(tasks, DropFixedSuperTableTask) # TODO: dicy
1179 1180 1181 1182
        # elif ( newState.equals(AnyState.STATE_TABLE_ONLY) ): # data deleted
            # self.assertNoTask(tasks, TaskDropDb)
            # self.assertNoTask(tasks, TaskDropSuperTable)
            # self.assertNoTask(tasks, TaskAddData)
S
Steven Li 已提交
1183
            # self.hasSuccess(tasks, DeleteDataTasks)
S
Shuduo Sang 已提交
1184 1185
        else:  # should be STATE_HAS_DATA
            if (not self.hasTask(tasks, TaskCreateDb)
1186
                ):  # only if we didn't create one
S
Shuduo Sang 已提交
1187 1188 1189
                # we shouldn't have dropped it
                self.assertNoTask(tasks, TaskDropDb)
            if (not self.hasTask(tasks, TaskCreateSuperTable)
1190
                    ):  # if we didn't create the table
S
Shuduo Sang 已提交
1191 1192
                # we should not have a task that drops it
                self.assertNoTask(tasks, TaskDropSuperTable)
1193
            # self.assertIfExistThenSuccess(tasks, ReadFixedDataTask)
S
Steven Li 已提交
1194

S
Shuduo Sang 已提交
1195

1196
class StateMechine:
1197 1198 1199
    def __init__(self, db: Database): 
        self._db = db
        # transitition target probabilities, indexed with value of STATE_EMPTY, STATE_DB_ONLY, etc.
1200
        self._stateWeights = [1, 2, 10, 40]
S
Shuduo Sang 已提交
1201

1202 1203
    def init(self, dbc: DbConn): # late initailization, don't save the dbConn
        self._curState = self._findCurrentState(dbc)  # starting state
1204
        Logging.debug("Found Starting State: {}".format(self._curState))
1205 1206

    # TODO: seems no lnoger used, remove?
1207 1208 1209
    def getCurrentState(self):
        return self._curState

1210 1211 1212
    def hasDatabase(self):
        return self._curState.canDropDb()  # ha, can drop DB means it has one

1213
    # May be slow, use cautionsly...
S
Shuduo Sang 已提交
1214
    def getTaskTypes(self):  # those that can run (directly/indirectly) from the current state
1215 1216 1217 1218 1219 1220
        def typesToStrings(types):
            ss = []
            for t in types:
                ss.append(t.__name__)
            return ss

S
Shuduo Sang 已提交
1221
        allTaskClasses = StateTransitionTask.__subclasses__()  # all state transition tasks
1222 1223
        firstTaskTypes = []
        for tc in allTaskClasses:
S
Shuduo Sang 已提交
1224
            # t = tc(self) # create task object
1225 1226
            if tc.canBeginFrom(self._curState):
                firstTaskTypes.append(tc)
S
Shuduo Sang 已提交
1227 1228 1229 1230 1231 1232 1233 1234
        # now we have all the tasks that can begin directly from the current
        # state, let's figure out the INDIRECT ones
        taskTypes = firstTaskTypes.copy()  # have to have these
        for task1 in firstTaskTypes:  # each task type gathered so far
            endState = task1.getEndState()  # figure the end state
            if endState is None:  # does not change end state
                continue  # no use, do nothing
            for tc in allTaskClasses:  # what task can further begin from there?
1235
                if tc.canBeginFrom(endState) and (tc not in firstTaskTypes):
S
Shuduo Sang 已提交
1236
                    taskTypes.append(tc)  # gather it
1237 1238

        if len(taskTypes) <= 0:
S
Shuduo Sang 已提交
1239 1240 1241
            raise RuntimeError(
                "No suitable task types found for state: {}".format(
                    self._curState))
1242
        Logging.debug(
S
Shuduo Sang 已提交
1243 1244 1245
            "[OPS] Tasks found for state {}: {}".format(
                self._curState,
                typesToStrings(taskTypes)))
1246 1247
        return taskTypes

1248
    def _findCurrentState(self, dbc: DbConn):
S
Shuduo Sang 已提交
1249
        ts = time.time()  # we use this to debug how fast/slow it is to do the various queries to find the current DB state
1250 1251
        dbName =self._db.getName()
        if not dbc.existsDatabase(dbName): # dbc.hasDatabases():  # no database?!
1252
            Logging.debug( "[STT] empty database found, between {} and {}".format(ts, time.time()))
1253
            return StateEmpty()
S
Shuduo Sang 已提交
1254 1255
        # did not do this when openning connection, and this is NOT the worker
        # thread, which does this on their own
1256
        dbc.use(dbName)
1257
        if not dbc.hasTables():  # no tables
1258
            Logging.debug("[STT] DB_ONLY found, between {} and {}".format(ts, time.time()))
1259
            return StateDbOnly()
1260

1261 1262
        sTable = self._db.getFixedSuperTable()
        if sTable.hasRegTables(dbc, dbName):  # no regular tables
1263
            Logging.debug("[STT] SUPER_TABLE_ONLY found, between {} and {}".format(ts, time.time()))
1264
            return StateSuperTableOnly()
S
Shuduo Sang 已提交
1265
        else:  # has actual tables
1266
            Logging.debug("[STT] HAS_DATA found, between {} and {}".format(ts, time.time()))
1267 1268
            return StateHasData()

1269 1270
    # We transition the system to a new state by examining the current state itself
    def transition(self, tasks, dbc: DbConn):
S
Shuduo Sang 已提交
1271
        if (len(tasks) == 0):  # before 1st step, or otherwise empty
1272
            Logging.debug("[STT] Starting State: {}".format(self._curState))
S
Shuduo Sang 已提交
1273
            return  # do nothing
1274

S
Shuduo Sang 已提交
1275
        # this should show up in the server log, separating steps
1276
        dbc.execute("show dnodes")
1277 1278 1279 1280

        # Generic Checks, first based on the start state
        if self._curState.canCreateDb():
            self._curState.assertIfExistThenSuccess(tasks, TaskCreateDb)
S
Shuduo Sang 已提交
1281 1282
            # self.assertAtMostOneSuccess(tasks, CreateDbTask) # not really, in
            # case of multiple creation and drops
1283 1284

        if self._curState.canDropDb():
1285
            if gSvcMgr == None: # only if we are running as client-only
S
Steven Li 已提交
1286
                self._curState.assertIfExistThenSuccess(tasks, TaskDropDb)
S
Shuduo Sang 已提交
1287 1288
            # self.assertAtMostOneSuccess(tasks, DropDbTask) # not really in
            # case of drop-create-drop
1289 1290 1291

        # if self._state.canCreateFixedTable():
            # self.assertIfExistThenSuccess(tasks, CreateFixedTableTask) # Not true, DB may be dropped
S
Shuduo Sang 已提交
1292 1293
            # self.assertAtMostOneSuccess(tasks, CreateFixedTableTask) # not
            # really, in case of create-drop-create
1294 1295 1296

        # if self._state.canDropFixedTable():
            # self.assertIfExistThenSuccess(tasks, DropFixedTableTask) # Not True, the whole DB may be dropped
S
Shuduo Sang 已提交
1297 1298
            # self.assertAtMostOneSuccess(tasks, DropFixedTableTask) # not
            # really in case of drop-create-drop
1299 1300

        # if self._state.canAddData():
S
Shuduo Sang 已提交
1301 1302
        # self.assertIfExistThenSuccess(tasks, AddFixedDataTask)  # not true
        # actually
1303 1304 1305 1306

        # if self._state.canReadData():
            # Nothing for sure

1307
        newState = self._findCurrentState(dbc)
1308
        Logging.debug("[STT] New DB state determined: {}".format(newState))
S
Shuduo Sang 已提交
1309 1310
        # can old state move to new state through the tasks?
        self._curState.verifyTasksToState(tasks, newState)
1311 1312 1313
        self._curState = newState

    def pickTaskType(self):
S
Shuduo Sang 已提交
1314 1315
        # all the task types we can choose from at curent state
        taskTypes = self.getTaskTypes()
1316 1317 1318
        weights = []
        for tt in taskTypes:
            endState = tt.getEndState()
S
Shuduo Sang 已提交
1319 1320 1321
            if endState is not None:
                # TODO: change to a method
                weights.append(self._stateWeights[endState.getValIndex()])
1322
            else:
S
Shuduo Sang 已提交
1323 1324
                # read data task, default to 10: TODO: change to a constant
                weights.append(10)
1325
        i = self._weighted_choice_sub(weights)
1326
        # Logging.debug(" (weighted random:{}/{}) ".format(i, len(taskTypes)))
1327 1328
        return taskTypes[i]

S
Shuduo Sang 已提交
1329 1330 1331 1332 1333
    # ref:
    # https://eli.thegreenplace.net/2010/01/22/weighted-random-generation-in-python/
    def _weighted_choice_sub(self, weights):
        # TODO: use our dice to ensure it being determinstic?
        rnd = random.random() * sum(weights)
1334 1335 1336 1337
        for i, w in enumerate(weights):
            rnd -= w
            if rnd < 0:
                return i
1338

1339 1340 1341 1342 1343 1344
class Database:
    ''' We use this to represent an actual TDengine database inside a service instance,
        possibly in a cluster environment.

        For now we use it to manage state transitions in that database
    '''
1345 1346 1347 1348 1349
    _clsLock = threading.Lock() # class wide lock
    _lastInt = 101  # next one is initial integer
    _lastTick = 0
    _lastLaggingTick = 0 # lagging tick, for unsequenced insersions

1350 1351 1352 1353
    def __init__(self, dbNum: int, dbc: DbConn): # TODO: remove dbc
        self._dbNum = dbNum # we assign a number to databases, for our testing purpose
        self._stateMachine = StateMechine(self)
        self._stateMachine.init(dbc)
1354
          
1355
        self._lock = threading.RLock()
S
Shuduo Sang 已提交
1356

1357 1358
    def getStateMachine(self) -> StateMechine:
        return self._stateMachine
S
Shuduo Sang 已提交
1359

1360 1361
    def getDbNum(self):
        return self._dbNum
1362

1363 1364
    def getName(self):
        return "db_{}".format(self._dbNum)
1365

1366 1367 1368 1369 1370 1371
    def filterTasks(self, inTasks: List[Task]): # Pick out those belonging to us
        outTasks = []
        for task in inTasks:
            if task.getDb().isSame(self):
                outTasks.append(task)
        return outTasks
1372

1373 1374 1375 1376 1377
    def isSame(self, other):
        return self._dbNum == other._dbNum

    def exists(self, dbc: DbConn):
        return dbc.existsDatabase(self.getName())
1378

1379 1380 1381 1382 1383 1384 1385
    @classmethod
    def getFixedSuperTableName(cls):
        return "fs_table"

    @classmethod
    def getFixedSuperTable(cls) -> TdSuperTable:
        return TdSuperTable(cls.getFixedSuperTableName())
1386 1387 1388 1389 1390 1391

    # We aim to create a starting time tick, such that, whenever we run our test here once
    # We should be able to safely create 100,000 records, which will not have any repeated time stamp
    # when we re-run the test in 3 minutes (180 seconds), basically we should expand time duration
    # by a factor of 500.
    # TODO: what if it goes beyond 10 years into the future
1392
    # TODO: fix the error as result of above: "tsdb timestamp is out of range"
1393 1394
    @classmethod
    def setupLastTick(cls):
1395
        t1 = datetime.datetime(2020, 6, 1)
1396
        t2 = datetime.datetime.now()
S
Shuduo Sang 已提交
1397 1398 1399 1400
        # maybe a very large number, takes 69 years to exceed Python int range
        elSec = int(t2.timestamp() - t1.timestamp())
        elSec2 = (elSec % (8 * 12 * 30 * 24 * 60 * 60 / 500)) * \
            500  # a number representing seconds within 10 years
1401
        # print("elSec = {}".format(elSec))
S
Shuduo Sang 已提交
1402 1403 1404
        t3 = datetime.datetime(2012, 1, 1)  # default "keep" is 10 years
        t4 = datetime.datetime.fromtimestamp(
            t3.timestamp() + elSec2)  # see explanation above
1405
        Logging.info("Setting up TICKS to start from: {}".format(t4))
1406 1407
        return t4

1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419
    @classmethod
    def getNextTick(cls):        
        with cls._clsLock:  # prevent duplicate tick
            if cls._lastLaggingTick==0:
                # 10k at 1/20 chance, should be enough to avoid overlaps
                cls._lastLaggingTick = cls.setupLastTick() + datetime.timedelta(0, -10000)                 
            if cls._lastTick==0: # should be quite a bit into the future
                cls._lastTick = cls.setupLastTick()  

            if Dice.throw(20) == 0:  # 1 in 20 chance, return lagging tick
                cls._lastLaggingTick += datetime.timedelta(0, 1) # Go back in time 100 seconds
                return cls._lastLaggingTick 
S
Shuduo Sang 已提交
1420 1421
            else:  # regular
                # add one second to it
1422 1423
                cls._lastTick += datetime.timedelta(0, 1)
                return cls._lastTick
1424 1425

    def getNextInt(self):
1426 1427 1428
        with self._lock:
            self._lastInt += 1
            return self._lastInt
1429 1430

    def getNextBinary(self):
S
Shuduo Sang 已提交
1431 1432
        return "Beijing_Shanghai_Los_Angeles_New_York_San_Francisco_Chicago_Beijing_Shanghai_Los_Angeles_New_York_San_Francisco_Chicago_{}".format(
            self.getNextInt())
1433 1434

    def getNextFloat(self):
1435 1436 1437
        ret = 0.9 + self.getNextInt()
        # print("Float obtained: {}".format(ret))
        return ret
S
Shuduo Sang 已提交
1438

1439 1440 1441 1442 1443 1444 1445 1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468 1469 1470 1471 1472 1473 1474 1475 1476 1477 1478 1479 1480 1481 1482 1483 1484 1485 1486 1487

class DbManager():
    ''' This is a wrapper around DbConn(), to make it easier to use. 

        TODO: rename this to DbConnManager
    '''
    def __init__(self):
        self.tableNumQueue = LinearQueue() # TODO: delete?
        # self.openDbServerConnection()
        self._dbConn = DbConn.createNative() if (
            gConfig.connector_type == 'native') else DbConn.createRest()
        try:
            self._dbConn.open()  # may throw taos.error.ProgrammingError: disconnected
        except taos.error.ProgrammingError as err:
            # print("Error type: {}, msg: {}, value: {}".format(type(err), err.msg, err))
            if (err.msg == 'client disconnected'):  # cannot open DB connection
                print(
                    "Cannot establish DB connection, please re-run script without parameter, and follow the instructions.")
                sys.exit(2)
            else:
                print("Failed to connect to DB, errno = {}, msg: {}"
                    .format(Helper.convertErrno(err.errno), err.msg))
                raise
        except BaseException:
            print("[=] Unexpected exception")
            raise

        # Do this after dbConn is in proper shape
        # Moved to Database()
        # self._stateMachine = StateMechine(self._dbConn)

    def getDbConn(self):
        return self._dbConn

    # TODO: not used any more, to delete
    def pickAndAllocateTable(self):  # pick any table, and "use" it
        return self.tableNumQueue.pickAndAllocate()

    # TODO: Not used any more, to delete
    def addTable(self):
        with self._lock:
            tIndex = self.tableNumQueue.push()
        return tIndex

    # Not used any more, to delete
    def releaseTable(self, i):  # return the table back, so others can use it
        self.tableNumQueue.release(i)    

    # TODO: not used any more, delete
S
Steven Li 已提交
1488
    def getTableNameToDelete(self):
S
Shuduo Sang 已提交
1489 1490
        tblNum = self.tableNumQueue.pop()  # TODO: race condition!
        if (not tblNum):  # maybe false
1491
            return False
S
Shuduo Sang 已提交
1492

S
Steven Li 已提交
1493 1494
        return "table_{}".format(tblNum)

1495
    def cleanUp(self):
S
Shuduo Sang 已提交
1496 1497
        self._dbConn.close()

1498
class TaskExecutor():
1499
    class BoundedList:
S
Shuduo Sang 已提交
1500
        def __init__(self, size=10):
1501 1502
            self._size = size
            self._list = []
S
Steven Li 已提交
1503
            self._lock = threading.Lock()
1504

S
Shuduo Sang 已提交
1505
        def add(self, n: int):
S
Steven Li 已提交
1506 1507 1508 1509 1510 1511 1512 1513 1514 1515 1516 1517 1518 1519 1520 1521 1522 1523 1524 1525 1526 1527 1528 1529 1530 1531
            with self._lock:
                if not self._list:  # empty
                    self._list.append(n)
                    return
                # now we should insert
                nItems = len(self._list)
                insPos = 0
                for i in range(nItems):
                    insPos = i
                    if n <= self._list[i]:  # smaller than this item, time to insert
                        break  # found the insertion point
                    insPos += 1  # insert to the right

                if insPos == 0:  # except for the 1st item, # TODO: elimiate first item as gating item
                    return  # do nothing

                # print("Inserting at postion {}, value: {}".format(insPos, n))
                self._list.insert(insPos, n)  # insert

                newLen = len(self._list)
                if newLen <= self._size:
                    return  # do nothing
                elif newLen == (self._size + 1):
                    del self._list[0]  # remove the first item
                else:
                    raise RuntimeError("Corrupt Bounded List")
1532 1533 1534 1535 1536 1537

        def __str__(self):
            return repr(self._list)

    _boundedList = BoundedList()

1538 1539 1540
    def __init__(self, curStep):
        self._curStep = curStep

1541 1542 1543 1544
    @classmethod
    def getBoundedList(cls):
        return cls._boundedList

1545 1546 1547
    def getCurStep(self):
        return self._curStep

S
Shuduo Sang 已提交
1548
    def execute(self, task: Task, wt: WorkerThread):  # execute a task on a thread
1549
        task.execute(wt)
1550

1551 1552 1553 1554
    def recordDataMark(self, n: int):
        # print("[{}]".format(n), end="", flush=True)
        self._boundedList.add(n)

1555
    # def logInfo(self, msg):
1556
    #     Logging.info("    T[{}.x]: ".format(self._curStep) + msg)
1557

1558
    # def logDebug(self, msg):
1559
    #     Logging.debug("    T[{}.x]: ".format(self._curStep) + msg)
1560

S
Shuduo Sang 已提交
1561

S
Steven Li 已提交
1562
class Task():
1563 1564 1565 1566
    ''' A generic "Task" to be executed. For now we decide that there is no
        need to embed a DB connection here, we use whatever the Worker Thread has
        instead. But a task is always associated with a DB
    '''
1567 1568 1569 1570
    taskSn = 100

    @classmethod
    def allocTaskNum(cls):
S
Shuduo Sang 已提交
1571
        Task.taskSn += 1  # IMPORTANT: cannot use cls.taskSn, since each sub class will have a copy
1572
        # Logging.debug("Allocating taskSN: {}".format(Task.taskSn))
S
Steven Li 已提交
1573
        return Task.taskSn
1574

1575
    def __init__(self, execStats: ExecutionStats, db: Database):
S
Shuduo Sang 已提交
1576
        self._workerThread = None
1577
        self._err = None # type: Exception
1578
        self._aborted = False
1579
        self._curStep = None
S
Shuduo Sang 已提交
1580
        self._numRows = None  # Number of rows affected
1581

S
Shuduo Sang 已提交
1582
        # Assign an incremental task serial number
1583
        self._taskNum = self.allocTaskNum()
1584
        # Logging.debug("Creating new task {}...".format(self._taskNum))
1585

1586
        self._execStats = execStats
1587
        self._db = db # A task is always associated/for a specific DB
1588

1589
    def isSuccess(self):
S
Shuduo Sang 已提交
1590
        return self._err is None
1591

1592 1593 1594
    def isAborted(self):
        return self._aborted

S
Shuduo Sang 已提交
1595
    def clone(self):  # TODO: why do we need this again?
1596
        newTask = self.__class__(self._execStats, self._db)
1597 1598
        return newTask

1599 1600 1601
    def getDb(self):
        return self._db

1602
    def logDebug(self, msg):
S
Shuduo Sang 已提交
1603 1604 1605
        self._workerThread.logDebug(
            "Step[{}.{}] {}".format(
                self._curStep, self._taskNum, msg))
1606 1607

    def logInfo(self, msg):
S
Shuduo Sang 已提交
1608 1609 1610
        self._workerThread.logInfo(
            "Step[{}.{}] {}".format(
                self._curStep, self._taskNum, msg))
1611

1612
    def _executeInternal(self, te: TaskExecutor, wt: WorkerThread):
S
Shuduo Sang 已提交
1613 1614 1615
        raise RuntimeError(
            "To be implemeted by child classes, class name: {}".format(
                self.__class__.__name__))
1616

1617 1618 1619
    def _isErrAcceptable(self, errno, msg):
        if errno in [
                0x05,  # TSDB_CODE_RPC_NOT_READY
1620
                0x0B,  # Unable to establish connection, more details in TD-1648
1621
                # 0x200, # invalid SQL, TODO: re-examine with TD-934
1622 1623
                0x217, # "db not selected", client side defined error code
                0x218, # "Table does not exist" client side defined error code
1624 1625 1626 1627 1628 1629 1630 1631 1632
                0x360, 0x362, 
                0x369, # tag already exists
                0x36A, 0x36B, 0x36D,
                0x381, 
                0x380, # "db not selected"
                0x383,
                0x386,  # DB is being dropped?!
                0x503,
                0x510,  # vnode not in ready state
1633
                0x14,   # db not ready, errno changed
1634
                0x600,  # Invalid table ID, why?
1635 1636 1637
                1000  # REST catch-all error
            ]: 
            return True # These are the ALWAYS-ACCEPTABLE ones
1638 1639
        elif (errno in [ 0x0B ]) and gConfig.auto_start_service:
            return True # We may get "network unavilable" when restarting service
1640 1641 1642
        elif errno == 0x200 : # invalid SQL, we need to div in a bit more
            if msg.find("invalid column name") != -1:
                return True 
1643 1644 1645 1646
            elif msg.find("tags number not matched") != -1: # mismatched tags after modification
                return True
            elif msg.find("duplicated column names") != -1: # also alter table tag issues
                return True
1647
        elif gSvcMgr and (not gSvcMgr.isStable()): # We are managing service, and ...
1648
            Logging.info("Ignoring error when service starting/stopping: errno = {}, msg = {}".format(errno, msg))
S
Steven Li 已提交
1649
            return True
1650 1651 1652 1653
        
        return False # Not an acceptable error


1654 1655
    def execute(self, wt: WorkerThread):
        wt.verifyThreadSelf()
S
Shuduo Sang 已提交
1656
        self._workerThread = wt  # type: ignore
1657 1658

        te = wt.getTaskExecutor()
1659
        self._curStep = te.getCurStep()
S
Shuduo Sang 已提交
1660 1661
        self.logDebug(
            "[-] executing task {}...".format(self.__class__.__name__))
1662

1663
        self._err = None # TODO: type hint mess up?
1664 1665
        self._execStats.beginTaskType(self.__class__.__name__)  # mark beginning
        errno2 = None
1666 1667 1668

        # Now pick a database, and stick with it for the duration of the task execution
        dbName = self._db.getName()
1669
        try:
S
Shuduo Sang 已提交
1670
            self._executeInternal(te, wt)  # TODO: no return value?
1671
        except taos.error.ProgrammingError as err:
1672
            errno2 = Helper.convertErrno(err.errno)
1673
            if (gConfig.continue_on_exception):  # user choose to continue
1674
                self.logDebug("[=] Continue after TAOS exception: errno=0x{:X}, msg: {}, SQL: {}".format(
1675
                        errno2, err, wt.getDbConn().getLastSql()))
1676
                self._err = err
1677 1678
            elif self._isErrAcceptable(errno2, err.__str__()):
                self.logDebug("[=] Acceptable Taos library exception: errno=0x{:X}, msg: {}, SQL: {}".format(
1679
                        errno2, err, wt.getDbConn().getLastSql()))
1680
                print("_", end="", flush=True)
S
Shuduo Sang 已提交
1681
                self._err = err
1682
            else: # not an acceptable error
1683 1684 1685
                errMsg = "[=] Unexpected Taos library exception ({}): errno=0x{:X}, msg: {}, SQL: {}".format(
                    self.__class__.__name__,
                    errno2, err, wt.getDbConn().getLastSql())
1686
                self.logDebug(errMsg)
S
Shuduo Sang 已提交
1687
                if gConfig.debug:
1688 1689
                    # raise # so that we see full stack
                    traceback.print_exc()
1690 1691
                print(
                    "\n\n----------------------------\nProgram ABORTED Due to Unexpected TAOS Error: \n\n{}\n".format(errMsg) +
1692 1693 1694 1695
                    "----------------------------\n")
                # sys.exit(-1)
                self._err = err
                self._aborted = True
S
Shuduo Sang 已提交
1696
        except Exception as e:
S
Steven Li 已提交
1697
            self.logInfo("Non-TAOS exception encountered")
S
Shuduo Sang 已提交
1698
            self._err = e
S
Steven Li 已提交
1699
            self._aborted = True
1700
            traceback.print_exc()
1701
        except BaseException as e:
1702
            self.logInfo("Python base exception encountered")
1703
            self._err = e
1704
            self._aborted = True
S
Steven Li 已提交
1705
            traceback.print_exc()
1706
        except BaseException: # TODO: what is this again??!!
S
Shuduo Sang 已提交
1707 1708
            self.logDebug(
                "[=] Unexpected exception, SQL: {}".format(
1709
                    wt.getDbConn().getLastSql()))
1710
            raise
1711
        self._execStats.endTaskType(self.__class__.__name__, self.isSuccess())
S
Shuduo Sang 已提交
1712 1713 1714 1715

        self.logDebug("[X] task execution completed, {}, status: {}".format(
            self.__class__.__name__, "Success" if self.isSuccess() else "Failure"))
        # TODO: merge with above.
1716
        self._execStats.incExecCount(self.__class__.__name__, self.isSuccess(), errno2)
S
Steven Li 已提交
1717

1718
    # TODO: refactor away, just provide the dbConn
S
Shuduo Sang 已提交
1719
    def execWtSql(self, wt: WorkerThread, sql):  # execute an SQL on the worker thread
1720
        """ Haha """
1721 1722
        return wt.execSql(sql)

S
Shuduo Sang 已提交
1723
    def queryWtSql(self, wt: WorkerThread, sql):  # execute an SQL on the worker thread
1724 1725
        return wt.querySql(sql)

S
Shuduo Sang 已提交
1726
    def getQueryResult(self, wt: WorkerThread):  # execute an SQL on the worker thread
1727 1728 1729
        return wt.getQueryResult()


1730
class ExecutionStats:
1731
    def __init__(self):
S
Shuduo Sang 已提交
1732 1733
        # total/success times for a task
        self._execTimes: Dict[str, [int, int]] = {}
1734 1735 1736
        self._tasksInProgress = 0
        self._lock = threading.Lock()
        self._firstTaskStartTime = None
1737
        self._execStartTime = None
1738
        self._errors = {}
S
Shuduo Sang 已提交
1739 1740
        self._elapsedTime = 0.0  # total elapsed time
        self._accRunTime = 0.0  # accumulated run time
1741

1742 1743 1744
        self._failed = False
        self._failureReason = None

S
Steven Li 已提交
1745
    def __str__(self):
S
Shuduo Sang 已提交
1746 1747
        return "[ExecStats: _failed={}, _failureReason={}".format(
            self._failed, self._failureReason)
S
Steven Li 已提交
1748 1749

    def isFailed(self):
S
Shuduo Sang 已提交
1750
        return self._failed
S
Steven Li 已提交
1751

1752 1753 1754 1755 1756 1757
    def startExec(self):
        self._execStartTime = time.time()

    def endExec(self):
        self._elapsedTime = time.time() - self._execStartTime

1758
    def incExecCount(self, klassName, isSuccess, eno=None):  # TODO: add a lock here
1759 1760
        if klassName not in self._execTimes:
            self._execTimes[klassName] = [0, 0]
S
Shuduo Sang 已提交
1761 1762
        t = self._execTimes[klassName]  # tuple for the data
        t[0] += 1  # index 0 has the "total" execution times
1763
        if isSuccess:
S
Shuduo Sang 已提交
1764
            t[1] += 1  # index 1 has the "success" execution times
1765 1766 1767 1768 1769
        if eno != None:             
            if klassName not in self._errors:
                self._errors[klassName] = {}
            errors = self._errors[klassName]
            errors[eno] = errors[eno]+1 if eno in errors else 1
1770 1771 1772

    def beginTaskType(self, klassName):
        with self._lock:
S
Shuduo Sang 已提交
1773 1774
            if self._tasksInProgress == 0:  # starting a new round
                self._firstTaskStartTime = time.time()  # I am now the first task
1775 1776 1777 1778 1779
            self._tasksInProgress += 1

    def endTaskType(self, klassName, isSuccess):
        with self._lock:
            self._tasksInProgress -= 1
S
Shuduo Sang 已提交
1780
            if self._tasksInProgress == 0:  # all tasks have stopped
1781 1782 1783
                self._accRunTime += (time.time() - self._firstTaskStartTime)
                self._firstTaskStartTime = None

1784 1785 1786 1787
    def registerFailure(self, reason):
        self._failed = True
        self._failureReason = reason

1788
    def printStats(self):
1789
        Logging.info(
S
Shuduo Sang 已提交
1790
            "----------------------------------------------------------------------")
1791
        Logging.info(
S
Shuduo Sang 已提交
1792 1793 1794
            "| Crash_Gen test {}, with the following stats:". format(
                "FAILED (reason: {})".format(
                    self._failureReason) if self._failed else "SUCCEEDED"))
1795
        Logging.info("| Task Execution Times (success/total):")
1796
        execTimesAny = 0.001 # avoid div by zero
S
Shuduo Sang 已提交
1797
        for k, n in self._execTimes.items():
1798
            execTimesAny += n[0]
1799 1800 1801 1802 1803 1804 1805
            errStr = None
            if k in self._errors:
                errors = self._errors[k]
                # print("errors = {}".format(errors))
                errStrs = ["0x{:X}:{}".format(eno, n) for (eno, n) in errors.items()]
                # print("error strings = {}".format(errStrs))
                errStr = ", ".join(errStrs) 
1806
            Logging.info("|    {0:<24}: {1}/{2} (Errors: {3})".format(k, n[1], n[0], errStr))
S
Shuduo Sang 已提交
1807

1808
        Logging.info(
S
Shuduo Sang 已提交
1809
            "| Total Tasks Executed (success or not): {} ".format(execTimesAny))
1810
        Logging.info(
S
Shuduo Sang 已提交
1811 1812
            "| Total Tasks In Progress at End: {}".format(
                self._tasksInProgress))
1813
        Logging.info(
S
Shuduo Sang 已提交
1814 1815
            "| Total Task Busy Time (elapsed time when any task is in progress): {:.3f} seconds".format(
                self._accRunTime))
1816
        Logging.info(
S
Shuduo Sang 已提交
1817
            "| Average Per-Task Execution Time: {:.3f} seconds".format(self._accRunTime / execTimesAny))
1818
        Logging.info(
S
Shuduo Sang 已提交
1819 1820
            "| Total Elapsed Time (from wall clock): {:.3f} seconds".format(
                self._elapsedTime))
1821 1822 1823
        Logging.info("| Top numbers written: {}".format(TaskExecutor.getBoundedList()))
        Logging.info("| Active DB Native Connections (now): {}".format(DbConnNative.totalConnections))
        Logging.info("| Longest native query time: {:.3f} seconds, started: {}".
1824 1825
            format(MyTDSql.longestQueryTime, 
                time.strftime("%x %X", time.localtime(MyTDSql.lqStartTime))) )
1826 1827
        Logging.info("| Longest native query: {}".format(MyTDSql.longestQuery))
        Logging.info(
S
Shuduo Sang 已提交
1828
            "----------------------------------------------------------------------")
1829 1830 1831


class StateTransitionTask(Task):
1832 1833 1834 1835 1836
    LARGE_NUMBER_OF_TABLES = 35
    SMALL_NUMBER_OF_TABLES = 3
    LARGE_NUMBER_OF_RECORDS = 50
    SMALL_NUMBER_OF_RECORDS = 3

1837
    @classmethod
S
Shuduo Sang 已提交
1838
    def getInfo(cls):  # each sub class should supply their own information
1839 1840
        raise RuntimeError("Overriding method expected")

S
Shuduo Sang 已提交
1841
    _endState = None
1842
    @classmethod
S
Shuduo Sang 已提交
1843
    def getEndState(cls):  # TODO: optimize by calling it fewer times
1844 1845
        raise RuntimeError("Overriding method expected")

1846 1847 1848
    # @classmethod
    # def getBeginStates(cls):
    #     return cls.getInfo()[0]
1849

1850 1851 1852
    # @classmethod
    # def getEndState(cls): # returning the class name
    #     return cls.getInfo()[0]
1853 1854

    @classmethod
1855 1856 1857
    def canBeginFrom(cls, state: AnyState):
        # return state.getValue() in cls.getBeginStates()
        raise RuntimeError("must be overriden")
1858

1859 1860
    @classmethod
    def getRegTableName(cls, i):
1861
        return "reg_table_{}".format(i)
1862

1863 1864
    def execute(self, wt: WorkerThread):
        super().execute(wt)
S
Shuduo Sang 已提交
1865 1866


1867
class TaskCreateDb(StateTransitionTask):
1868
    @classmethod
1869
    def getEndState(cls):
S
Shuduo Sang 已提交
1870
        return StateDbOnly()
1871

1872 1873 1874 1875
    @classmethod
    def canBeginFrom(cls, state: AnyState):
        return state.canCreateDb()

1876
    # Actually creating the database(es)
1877
    def _executeInternal(self, te: TaskExecutor, wt: WorkerThread):
1878
        # was: self.execWtSql(wt, "create database db")
1879 1880 1881 1882 1883 1884
        repStr = ""
        if gConfig.max_replicas != 1:
            numReplica = Dice.throw(gConfig.max_replicas) + 1 # 1,2 ... N
            repStr = "replica {}".format(numReplica)
        self.execWtSql(wt, "create database {} {}"
            .format(self._db.getName(), repStr) )
1885

1886
class TaskDropDb(StateTransitionTask):
1887
    @classmethod
1888 1889
    def getEndState(cls):
        return StateEmpty()
1890

1891 1892 1893 1894
    @classmethod
    def canBeginFrom(cls, state: AnyState):
        return state.canDropDb()

1895
    def _executeInternal(self, te: TaskExecutor, wt: WorkerThread):
1896
        self.execWtSql(wt, "drop database {}".format(self._db.getName()))
1897
        Logging.debug("[OPS] database dropped at {}".format(time.time()))
1898

1899
class TaskCreateSuperTable(StateTransitionTask):
1900
    @classmethod
1901 1902
    def getEndState(cls):
        return StateSuperTableOnly()
1903

1904 1905
    @classmethod
    def canBeginFrom(cls, state: AnyState):
1906
        return state.canCreateFixedSuperTable()
1907

1908
    def _executeInternal(self, te: TaskExecutor, wt: WorkerThread):
1909
        if not self._db.exists(wt.getDbConn()):
1910
            Logging.debug("Skipping task, no DB yet")
1911 1912
            return

1913
        sTable = self._db.getFixedSuperTable() # type: TdSuperTable
1914
        # wt.execSql("use db")    # should always be in place
1915 1916
        sTable.create(wt.getDbConn(), self._db.getName(), 
            {'ts':'timestamp', 'speed':'int'}, {'b':'binary(200)', 'f':'float'})
1917
        # self.execWtSql(wt,"create table db.{} (ts timestamp, speed int) tags (b binary(200), f float) ".format(tblName))
S
Shuduo Sang 已提交
1918 1919
        # No need to create the regular tables, INSERT will do that
        # automatically
1920

S
Steven Li 已提交
1921

1922 1923 1924 1925
class TdSuperTable:
    def __init__(self, stName):
        self._stName = stName

1926 1927 1928
    def getName(self):
        return self._stName

1929 1930 1931 1932 1933
    # TODO: odd semantic, create() method is usually static?
    def create(self, dbc, dbName, cols: dict, tags: dict):
        '''Creating a super table'''
        sql = "CREATE TABLE {}.{} ({}) TAGS ({})".format(
            dbName,
1934 1935 1936 1937 1938 1939
            self._stName,
            ",".join(['%s %s'%(k,v) for (k,v) in cols.items()]),
            ",".join(['%s %s'%(k,v) for (k,v) in tags.items()])
            )
        dbc.execute(sql)        

1940
    def getRegTables(self, dbc: DbConn, dbName: str):
1941
        try:
1942
            dbc.query("select TBNAME from {}.{}".format(dbName, self._stName))  # TODO: analyze result set later            
1943
        except taos.error.ProgrammingError as err:                    
1944
            errno2 = Helper.convertErrno(err.errno) 
1945
            Logging.debug("[=] Failed to get tables from super table: errno=0x{:X}, msg: {}".format(errno2, err))
1946 1947 1948 1949 1950
            raise

        qr = dbc.getQueryResult()
        return [v[0] for v in qr] # list transformation, ref: https://stackoverflow.com/questions/643823/python-list-transformation

1951 1952
    def hasRegTables(self, dbc: DbConn, dbName: str):
        return dbc.query("SELECT * FROM {}.{}".format(dbName, self._stName)) > 0
1953

1954 1955
    def ensureTable(self, dbc: DbConn, dbName: str, regTableName: str):
        sql = "select tbname from {}.{} where tbname in ('{}')".format(dbName, self._stName, regTableName)
1956 1957
        if dbc.query(sql) >= 1 : # reg table exists already
            return
1958 1959
        sql = "CREATE TABLE {}.{} USING {}.{} tags ({})".format(
            dbName, regTableName, dbName, self._stName, self._getTagStrForSql(dbc, dbName)
1960 1961 1962
        )
        dbc.execute(sql)

1963 1964
    def _getTagStrForSql(self, dbc, dbName: str) :
        tags = self._getTags(dbc, dbName)
1965 1966 1967 1968 1969 1970 1971 1972 1973 1974 1975 1976 1977
        tagStrs = []
        for tagName in tags: 
            tagType = tags[tagName]
            if tagType == 'BINARY':
                tagStrs.append("'Beijing-Shanghai-LosAngeles'")
            elif tagType == 'FLOAT':
                tagStrs.append('9.9')
            elif tagType == 'INT':
                tagStrs.append('88')
            else:
                raise RuntimeError("Unexpected tag type: {}".format(tagType))
        return ", ".join(tagStrs)

1978 1979
    def _getTags(self, dbc, dbName) -> dict:
        dbc.query("DESCRIBE {}.{}".format(dbName, self._stName))
1980 1981 1982 1983 1984 1985
        stCols = dbc.getQueryResult()
        # print(stCols)
        ret = {row[0]:row[1] for row in stCols if row[3]=='TAG'} # name:type
        # print("Tags retrieved: {}".format(ret))
        return ret

1986 1987
    def addTag(self, dbc, dbName, tagName, tagType):
        if tagName in self._getTags(dbc, dbName): # already 
1988 1989
            return
        # sTable.addTag("extraTag", "int")
1990
        sql = "alter table {}.{} add tag {} {}".format(dbName, self._stName, tagName, tagType)
1991 1992
        dbc.execute(sql)

1993 1994
    def dropTag(self, dbc, dbName, tagName):
        if not tagName in self._getTags(dbc, dbName): # don't have this tag
1995
            return
1996
        sql = "alter table {}.{} drop tag {}".format(dbName, self._stName, tagName)
1997 1998
        dbc.execute(sql)

1999 2000
    def changeTag(self, dbc, dbName, oldTag, newTag):
        tags = self._getTags(dbc, dbName)
2001 2002 2003 2004
        if not oldTag in tags: # don't have this tag
            return
        if newTag in tags: # already have this tag
            return
2005
        sql = "alter table {}.{} change tag {} {}".format(dbName, self._stName, oldTag, newTag)
2006 2007
        dbc.execute(sql)

2008
class TaskReadData(StateTransitionTask):
2009
    @classmethod
2010
    def getEndState(cls):
S
Shuduo Sang 已提交
2011
        return None  # meaning doesn't affect state
2012

2013 2014 2015 2016
    @classmethod
    def canBeginFrom(cls, state: AnyState):
        return state.canReadData()

2017
    def _executeInternal(self, te: TaskExecutor, wt: WorkerThread):
2018
        sTable = self._db.getFixedSuperTable()
2019

2020 2021
        # 1 in 5 chance, simulate a broken connection. 
        if random.randrange(5) == 0:  # TODO: break connection in all situations
2022 2023
            wt.getDbConn().close()
            wt.getDbConn().open()
2024
            print("_r", end="", flush=True)
2025
        
2026
        dbc = wt.getDbConn()
2027 2028
        dbName = self._db.getName()
        for rTbName in sTable.getRegTables(dbc, dbName):  # regular tables
2029
            aggExpr = Dice.choice([
2030 2031 2032
                '*',
                'count(*)',
                'avg(speed)',
2033
                # 'twa(speed)', # TODO: this one REQUIRES a where statement, not reasonable
2034 2035
                'sum(speed)', 
                'stddev(speed)', 
2036
                # SELECTOR functions
2037 2038 2039
                'min(speed)', 
                'max(speed)', 
                'first(speed)', 
2040
                'last(speed)',
2041 2042 2043
                'top(speed, 50)', # TODO: not supported?
                'bottom(speed, 50)', # TODO: not supported?
                'apercentile(speed, 10)', # TODO: TD-1316
2044 2045 2046 2047 2048
                'last_row(speed)',
                # Transformation Functions
                # 'diff(speed)', # TODO: no supported?!
                'spread(speed)'
                ]) # TODO: add more from 'top'
2049 2050 2051
            filterExpr = Dice.choice([ # TODO: add various kind of WHERE conditions
                None
            ])
2052
            try:
2053
                # Run the query against the regular table first
2054
                dbc.execute("select {} from {}.{}".format(aggExpr, dbName, rTbName))
2055
                # Then run it against the super table
2056
                if aggExpr not in ['stddev(speed)']: #TODO: STDDEV not valid for super tables?!
2057
                    dbc.execute("select {} from {}.{}".format(aggExpr, dbName, sTable.getName()))
2058
            except taos.error.ProgrammingError as err:                    
2059
                errno2 = Helper.convertErrno(err.errno)
2060
                Logging.debug("[=] Read Failure: errno=0x{:X}, msg: {}, SQL: {}".format(errno2, err, dbc.getLastSql()))
2061
                raise
S
Shuduo Sang 已提交
2062

2063
class TaskDropSuperTable(StateTransitionTask):
2064
    @classmethod
2065
    def getEndState(cls):
S
Shuduo Sang 已提交
2066
        return StateDbOnly()
2067

2068 2069
    @classmethod
    def canBeginFrom(cls, state: AnyState):
2070
        return state.canDropFixedSuperTable()
2071

2072
    def _executeInternal(self, te: TaskExecutor, wt: WorkerThread):
2073
        # 1/2 chance, we'll drop the regular tables one by one, in a randomized sequence
S
Shuduo Sang 已提交
2074
        if Dice.throw(2) == 0:
2075
            # print("_7_", end="", flush=True)
S
Shuduo Sang 已提交
2076 2077 2078 2079
            tblSeq = list(range(
                2 + (self.LARGE_NUMBER_OF_TABLES if gConfig.larger_data else self.SMALL_NUMBER_OF_TABLES)))
            random.shuffle(tblSeq)
            tickOutput = False  # if we have spitted out a "d" character for "drop regular table"
2080
            isSuccess = True
S
Shuduo Sang 已提交
2081
            for i in tblSeq:
2082
                regTableName = self.getRegTableName(i)  # "db.reg_table_{}".format(i)
2083
                try:
2084 2085
                    self.execWtSql(wt, "drop table {}.{}".
                        format(self._db.getName(), regTableName))  # nRows always 0, like MySQL
S
Shuduo Sang 已提交
2086
                except taos.error.ProgrammingError as err:
2087 2088
                    # correcting for strange error number scheme                    
                    errno2 = Helper.convertErrno(err.errno)
S
Shuduo Sang 已提交
2089
                    if (errno2 in [0x362]):  # mnode invalid table name
2090
                        isSuccess = False
2091
                        Logging.debug("[DB] Acceptable error when dropping a table")
S
Shuduo Sang 已提交
2092
                    continue  # try to delete next regular table
2093 2094

                if (not tickOutput):
S
Shuduo Sang 已提交
2095 2096
                    tickOutput = True  # Print only one time
                    if isSuccess:
2097 2098
                        print("d", end="", flush=True)
                    else:
S
Shuduo Sang 已提交
2099
                        print("f", end="", flush=True)
2100 2101

        # Drop the super table itself
2102 2103
        tblName = self._db.getFixedSuperTableName()
        self.execWtSql(wt, "drop table {}.{}".format(self._db.getName(), tblName))
2104

S
Shuduo Sang 已提交
2105

2106 2107 2108
class TaskAlterTags(StateTransitionTask):
    @classmethod
    def getEndState(cls):
S
Shuduo Sang 已提交
2109
        return None  # meaning doesn't affect state
2110 2111 2112

    @classmethod
    def canBeginFrom(cls, state: AnyState):
S
Shuduo Sang 已提交
2113
        return state.canDropFixedSuperTable()  # if we can drop it, we can alter tags
2114 2115

    def _executeInternal(self, te: TaskExecutor, wt: WorkerThread):
2116 2117
        # tblName = self._dbManager.getFixedSuperTableName()
        dbc = wt.getDbConn()
2118 2119
        sTable = self._db.getFixedSuperTable()
        dbName = self._db.getName()
2120
        dice = Dice.throw(4)
S
Shuduo Sang 已提交
2121
        if dice == 0:
2122
            sTable.addTag(dbc, dbName, "extraTag", "int")
2123
            # sql = "alter table db.{} add tag extraTag int".format(tblName)
S
Shuduo Sang 已提交
2124
        elif dice == 1:
2125
            sTable.dropTag(dbc, dbName, "extraTag")
2126
            # sql = "alter table db.{} drop tag extraTag".format(tblName)
S
Shuduo Sang 已提交
2127
        elif dice == 2:
2128
            sTable.dropTag(dbc, dbName, "newTag")
2129
            # sql = "alter table db.{} drop tag newTag".format(tblName)
S
Shuduo Sang 已提交
2130
        else:  # dice == 3
2131
            sTable.changeTag(dbc, dbName, "extraTag", "newTag")
2132 2133 2134 2135 2136 2137 2138 2139 2140 2141 2142 2143 2144 2145 2146 2147
            # sql = "alter table db.{} change tag extraTag newTag".format(tblName)

class TaskRestartService(StateTransitionTask):
    _isRunning = False
    _classLock = threading.Lock()

    @classmethod
    def getEndState(cls):
        return None  # meaning doesn't affect state

    @classmethod
    def canBeginFrom(cls, state: AnyState):
        if gConfig.auto_start_service:
            return state.canDropFixedSuperTable()  # Basicallly when we have the super table
        return False # don't run this otherwise

2148
    CHANCE_TO_RESTART_SERVICE = 200
2149 2150 2151 2152
    def _executeInternal(self, te: TaskExecutor, wt: WorkerThread):
        if not gConfig.auto_start_service: # only execute when we are in -a mode
            print("_a", end="", flush=True)
            return
2153

2154 2155 2156 2157 2158 2159
        with self._classLock:
            if self._isRunning:
                print("Skipping restart task, another running already")
                return
            self._isRunning = True

2160
        if Dice.throw(self.CHANCE_TO_RESTART_SERVICE) == 0: # 1 in N chance
2161 2162 2163
            dbc = wt.getDbConn()
            dbc.execute("show databases") # simple delay, align timing with other workers
            gSvcMgr.restart()
2164

2165
        self._isRunning = False
S
Shuduo Sang 已提交
2166

2167
class TaskAddData(StateTransitionTask):
S
Shuduo Sang 已提交
2168 2169
    # Track which table is being actively worked on
    activeTable: Set[int] = set()
2170

2171 2172 2173
    # We use these two files to record operations to DB, useful for power-off tests
    fAddLogReady = None # type: TextIOWrapper
    fAddLogDone  = None # type: TextIOWrapper
2174 2175 2176

    @classmethod
    def prepToRecordOps(cls):
S
Shuduo Sang 已提交
2177 2178
        if gConfig.record_ops:
            if (cls.fAddLogReady is None):
2179
                Logging.info(
S
Shuduo Sang 已提交
2180
                    "Recording in a file operations to be performed...")
2181
                cls.fAddLogReady = open("add_log_ready.txt", "w")
S
Shuduo Sang 已提交
2182
            if (cls.fAddLogDone is None):
2183
                Logging.info("Recording in a file operations completed...")
2184
                cls.fAddLogDone = open("add_log_done.txt", "w")
2185

2186
    @classmethod
2187 2188
    def getEndState(cls):
        return StateHasData()
2189 2190 2191 2192

    @classmethod
    def canBeginFrom(cls, state: AnyState):
        return state.canAddData()
S
Shuduo Sang 已提交
2193

2194
    def _executeInternal(self, te: TaskExecutor, wt: WorkerThread):
2195 2196
        # ds = self._dbManager # Quite DANGEROUS here, may result in multi-thread client access
        db = self._db
2197
        dbc = wt.getDbConn()
2198
        tblSeq = list(range(
S
Shuduo Sang 已提交
2199 2200 2201 2202
                self.LARGE_NUMBER_OF_TABLES if gConfig.larger_data else self.SMALL_NUMBER_OF_TABLES))
        random.shuffle(tblSeq)
        for i in tblSeq:
            if (i in self.activeTable):  # wow already active
2203
                print("x", end="", flush=True) # concurrent insertion
2204
            else:
S
Shuduo Sang 已提交
2205
                self.activeTable.add(i)  # marking it active
2206
            
2207
            sTable = db.getFixedSuperTable()
2208
            regTableName = self.getRegTableName(i)  # "db.reg_table_{}".format(i)
2209
            sTable.ensureTable(wt.getDbConn(), db.getName(), regTableName)  # Ensure the table exists           
2210 2211
           
            for j in range(self.LARGE_NUMBER_OF_RECORDS if gConfig.larger_data else self.SMALL_NUMBER_OF_RECORDS):  # number of records per table
2212
                nextInt = db.getNextInt()
2213
                nextTick = db.getNextTick()
2214 2215
                if gConfig.record_ops:
                    self.prepToRecordOps()
2216
                    self.fAddLogReady.write("Ready to write {} to {}\n".format(nextInt, regTableName))
2217 2218
                    self.fAddLogReady.flush()
                    os.fsync(self.fAddLogReady)
2219 2220
                sql = "insert into {}.{} values ('{}', {});".format( # removed: tags ('{}', {})
                    db.getName(),
S
Shuduo Sang 已提交
2221
                    regTableName,
2222 2223
                    # ds.getFixedSuperTableName(),
                    # ds.getNextBinary(), ds.getNextFloat(),
2224 2225
                    nextTick, nextInt)
                dbc.execute(sql)
S
Shuduo Sang 已提交
2226 2227
                # Successfully wrote the data into the DB, let's record it
                # somehow
2228
                te.recordDataMark(nextInt)
2229
                if gConfig.record_ops:
S
Shuduo Sang 已提交
2230 2231 2232
                    self.fAddLogDone.write(
                        "Wrote {} to {}\n".format(
                            nextInt, regTableName))
2233 2234
                    self.fAddLogDone.flush()
                    os.fsync(self.fAddLogDone)
2235 2236

                # Now read it back and verify, we might encounter an error if table is dropped
2237 2238 2239 2240 2241 2242 2243 2244 2245 2246 2247 2248 2249 2250 2251 2252 2253
                if gConfig.verify_data: # only if command line asks for it
                    try:
                        readBack = dbc.queryScalar("SELECT speed from {}.{} WHERE ts= '{}'".
                            format(db.getName(), regTableName, nextTick))
                        if readBack != nextInt :
                            raise taos.error.ProgrammingError(
                                "Failed to read back same data, wrote: {}, read: {}"
                                .format(nextInt, readBack), 0x999)
                    except taos.error.ProgrammingError as err:
                        errno = Helper.convertErrno(err.errno)
                        if errno in [0x991, 0x992]  : # not a single result
                            raise taos.error.ProgrammingError(
                                "Failed to read back same data for tick: {}, wrote: {}, read: {}"
                                .format(nextTick, nextInt, "Empty Result" if errno==0x991 else "Multiple Result"),
                                errno)
                        # Re-throw no matter what
                        raise
2254 2255
                

S
Shuduo Sang 已提交
2256
            self.activeTable.discard(i)  # not raising an error, unlike remove
2257 2258


2259

2260

2261

2262 2263 2264 2265 2266 2267 2268 2269 2270 2271 2272 2273 2274 2275 2276 2277 2278 2279 2280 2281 2282 2283 2284 2285 2286 2287 2288 2289
class ThreadStacks: # stack info for all threads
    def __init__(self):
        self._allStacks = {}
        allFrames = sys._current_frames()
        for th in threading.enumerate():                        
            stack = traceback.extract_stack(allFrames[th.ident])     
            self._allStacks[th.native_id] = stack

    def print(self, filteredEndName = None, filterInternal = False):
        for thNid, stack in self._allStacks.items(): # for each thread            
            lastFrame = stack[-1]
            if filteredEndName: # we need to filter out stacks that match this name                
                if lastFrame.name == filteredEndName : # end did not match
                    continue
            if filterInternal:
                if lastFrame.name in ['wait', 'invoke_excepthook', 
                    '_wait', # The Barrier exception
                    'svcOutputReader', # the svcMgr thread
                    '__init__']: # the thread that extracted the stack
                    continue # ignore
            # Now print
            print("\n<----- Thread Info for ID: {}".format(thNid))
            for frame in stack:
                # print(frame)
                print("File {filename}, line {lineno}, in {name}".format(
                    filename=frame.filename, lineno=frame.lineno, name=frame.name))
                print("    {}".format(frame.line))
            print("-----> End of Thread Info\n")
S
Shuduo Sang 已提交
2290

2291 2292 2293
class ClientManager:
    def __init__(self):
        print("Starting service manager")
2294 2295
        # signal.signal(signal.SIGTERM, self.sigIntHandler)
        # signal.signal(signal.SIGINT, self.sigIntHandler)
2296

2297
        self._status = Status.STATUS_RUNNING
2298 2299
        self.tc = None

2300 2301
        self.inSigHandler = False

2302
    def sigIntHandler(self, signalNumber, frame):
2303
        if self._status != Status.STATUS_RUNNING:
2304 2305 2306
            print("Repeated SIGINT received, forced exit...")
            # return  # do nothing if it's already not running
            sys.exit(-1)
2307
        self._status = Status.STATUS_STOPPING  # immediately set our status
2308

2309
        print("ClientManager: Terminating program...")
2310 2311
        self.tc.requestToStop()

2312 2313 2314 2315 2316 2317 2318 2319 2320 2321 2322 2323 2324 2325 2326 2327 2328 2329 2330 2331 2332 2333 2334 2335 2336 2337 2338 2339 2340 2341 2342 2343 2344 2345 2346 2347 2348 2349 2350 2351 2352
    def _doMenu(self):
        choice = ""
        while True:
            print("\nInterrupting Client Program, Choose an Action: ")
            print("1: Resume")
            print("2: Terminate")
            print("3: Show Threads")
            # Remember to update the if range below
            # print("Enter Choice: ", end="", flush=True)
            while choice == "":
                choice = input("Enter Choice: ")
                if choice != "":
                    break  # done with reading repeated input
            if choice in ["1", "2", "3"]:
                break  # we are done with whole method
            print("Invalid choice, please try again.")
            choice = ""  # reset
        return choice

    def sigUsrHandler(self, signalNumber, frame):
        print("Interrupting main thread execution upon SIGUSR1")
        if self.inSigHandler:  # already
            print("Ignoring repeated SIG_USR1...")
            return  # do nothing if it's already not running
        self.inSigHandler = True

        choice = self._doMenu()
        if choice == "1":
            print("Resuming execution...")
            time.sleep(1.0)
        elif choice == "2":
            print("Not implemented yet")
            time.sleep(1.0)
        elif choice == "3":
            ts = ThreadStacks()
            ts.print()
        else:
            raise RuntimeError("Invalid menu choice: {}".format(choice))

        self.inSigHandler = False

2353 2354 2355 2356 2357 2358 2359 2360 2361 2362 2363 2364 2365 2366 2367 2368 2369 2370 2371 2372 2373 2374 2375 2376 2377 2378 2379 2380 2381
    # TODO: need to revise how we verify data durability
    # def _printLastNumbers(self):  # to verify data durability
    #     dbManager = DbManager()
    #     dbc = dbManager.getDbConn()
    #     if dbc.query("show databases") <= 1:  # no database (we have a default called "log")
    #         return
    #     dbc.execute("use db")
    #     if dbc.query("show tables") == 0:  # no tables
    #         return

    #     sTbName = dbManager.getFixedSuperTableName()

    #     # get all regular tables
    #     # TODO: analyze result set later
    #     dbc.query("select TBNAME from db.{}".format(sTbName))
    #     rTables = dbc.getQueryResult()

    #     bList = TaskExecutor.BoundedList()
    #     for rTbName in rTables:  # regular tables
    #         dbc.query("select speed from db.{}".format(rTbName[0]))
    #         numbers = dbc.getQueryResult()
    #         for row in numbers:
    #             # print("<{}>".format(n), end="", flush=True)
    #             bList.add(row[0])

    #     print("Top numbers in DB right now: {}".format(bList))
    #     print("TDengine client execution is about to start in 2 seconds...")
    #     time.sleep(2.0)
    #     dbManager = None  # release?
2382

2383
    def run(self, svcMgr):    
2384
        # self._printLastNumbers()
2385
        global gConfig
2386

2387 2388 2389 2390
        # Prepare Tde Instance
        global gContainer
        tInst = gContainer.defTdeInstance = TdeInstance() # "subdir to hold the instance"

S
Shuduo Sang 已提交
2391
        dbManager = DbManager()  # Regular function
2392
        thPool = ThreadPool(gConfig.num_threads, gConfig.max_steps)
2393
        self.tc = ThreadCoordinator(thPool, dbManager)
2394
        
2395
        print("Starting client instance to: {}".format(tInst))
2396
        self.tc.run()
S
Steven Li 已提交
2397 2398
        # print("exec stats: {}".format(self.tc.getExecStats()))
        # print("TC failed = {}".format(self.tc.isFailed()))
2399
        if svcMgr: # gConfig.auto_start_service:
2400
            svcMgr.stopTaosService()
2401
            svcMgr = None
2402 2403
        # Print exec status, etc., AFTER showing messages from the server
        self.conclude()
S
Steven Li 已提交
2404
        # print("TC failed (2) = {}".format(self.tc.isFailed()))
S
Shuduo Sang 已提交
2405
        # Linux return code: ref https://shapeshed.com/unix-exit-codes/
2406 2407 2408 2409 2410 2411 2412 2413 2414 2415 2416 2417 2418 2419 2420 2421 2422 2423 2424
        ret = 1 if self.tc.isFailed() else 0
        self.tc.cleanup()

        # Release global variables
        gConfig = None
        gSvcMgr = None
        logger = None

        # Release variables here
        self.tc = None
        thPool = None
        dbManager = None

        gc.collect() # force garbage collection
        # h = hpy()
        # print("\n----- Final Python Heap -----\n")        
        # print(h.heap())

        return ret
2425 2426

    def conclude(self):
2427
        # self.tc.getDbManager().cleanUp() # clean up first, so we can show ZERO db connections
2428
        self.tc.printStats()
2429

2430
class MainExec:
2431 2432 2433
    def __init__(self):        
        self._clientMgr = None
        self._svcMgr = None
2434

2435 2436 2437
        signal.signal(signal.SIGTERM, self.sigIntHandler)
        signal.signal(signal.SIGINT,  self.sigIntHandler)
        signal.signal(signal.SIGUSR1, self.sigUsrHandler)  # different handler!
2438

2439 2440 2441 2442 2443 2444 2445 2446 2447 2448 2449 2450 2451
    def sigUsrHandler(self, signalNumber, frame):
        if self._clientMgr:
            self._clientMgr.sigUsrHandler(signalNumber, frame)
        elif self._svcMgr: # Only if no client mgr, we are running alone
            self._svcMgr.sigUsrHandler(signalNumber, frame)
        
    def sigIntHandler(self, signalNumber, frame):
        if self._svcMgr:
            self._svcMgr.sigIntHandler(signalNumber, frame)
        if self._clientMgr:
            self._clientMgr.sigIntHandler(signalNumber, frame)

    def runClient(self):
2452
        global gSvcMgr
2453
        if gConfig.auto_start_service:
2454
            self._svcMgr = ServiceManager()
2455
            gSvcMgr = self._svcMgr # hack alert
2456 2457 2458
            self._svcMgr.startTaosService() # we start, don't run
        
        self._clientMgr = ClientManager()
2459 2460 2461 2462
        ret = None
        try: 
            ret = self._clientMgr.run(self._svcMgr) # stop TAOS service inside
        except requests.exceptions.ConnectionError as err:
2463
            Logging.warning("Failed to open REST connection to DB: {}".format(err.getMessage()))
2464 2465
            # don't raise
        return ret
2466 2467

    def runService(self):
2468
        global gSvcMgr
2469
        self._svcMgr = ServiceManager()
2470 2471
        gSvcMgr = self._svcMgr # save it in a global variable TODO: hack alert

2472
        self._svcMgr.run() # run to some end state
2473 2474
        self._svcMgr = None 
        gSvcMgr = None        
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 2501 2502 2503 2504 2505 2506 2507 2508 2509 2510 2511 2512 2513 2514 2515 2516 2517 2518 2519 2520 2521 2522 2523 2524 2525 2526 2527 2528 2529 2530 2531 2532 2533 2534 2535 2536 2537 2538 2539 2540 2541 2542 2543 2544 2545 2546 2547 2548 2549 2550 2551 2552 2553 2554 2555 2556 2557 2558 2559 2560 2561 2562 2563 2564 2565 2566 2567 2568 2569
    def init(self): # TODO: refactor
        global gContainer
        gContainer = Container() # micky-mouse DI

        # Super cool Python argument library:
        # https://docs.python.org/3/library/argparse.html
        parser = argparse.ArgumentParser(
            formatter_class=argparse.RawDescriptionHelpFormatter,
            description=textwrap.dedent('''\
                TDengine Auto Crash Generator (PLEASE NOTICE the Prerequisites Below)
                ---------------------------------------------------------------------
                1. You build TDengine in the top level ./build directory, as described in offical docs
                2. You run the server there before this script: ./build/bin/taosd -c test/cfg

                '''))                      

        parser.add_argument(
            '-a',
            '--auto-start-service',
            action='store_true',
            help='Automatically start/stop the TDengine service (default: false)')
        parser.add_argument(
            '-b',
            '--max-dbs',
            action='store',
            default=0,
            type=int,
            help='Maximum number of DBs to keep, set to disable dropping DB. (default: 0)')
        parser.add_argument(
            '-c',
            '--connector-type',
            action='store',
            default='native',
            type=str,
            help='Connector type to use: native, rest, or mixed (default: 10)')
        parser.add_argument(
            '-d',
            '--debug',
            action='store_true',
            help='Turn on DEBUG mode for more logging (default: false)')
        parser.add_argument(
            '-e',
            '--run-tdengine',
            action='store_true',
            help='Run TDengine service in foreground (default: false)')
        parser.add_argument(
            '-i',
            '--max-replicas',
            action='store',
            default=1,
            type=int,
            help='Maximum number of replicas to use, when testing against clusters. (default: 1)')
        parser.add_argument(
            '-l',
            '--larger-data',
            action='store_true',
            help='Write larger amount of data during write operations (default: false)')
        parser.add_argument(
            '-p',
            '--per-thread-db-connection',
            action='store_true',
            help='Use a single shared db connection (default: false)')
        parser.add_argument(
            '-r',
            '--record-ops',
            action='store_true',
            help='Use a pair of always-fsynced fils to record operations performing + performed, for power-off tests (default: false)')
        parser.add_argument(
            '-s',
            '--max-steps',
            action='store',
            default=1000,
            type=int,
            help='Maximum number of steps to run (default: 100)')
        parser.add_argument(
            '-t',
            '--num-threads',
            action='store',
            default=5,
            type=int,
            help='Number of threads to run (default: 10)')
        parser.add_argument(
            '-v',
            '--verify-data',
            action='store_true',
            help='Verify data written in a number of places by reading back (default: false)')
        parser.add_argument(
            '-x',
            '--continue-on-exception',
            action='store_true',
            help='Continue execution after encountering unexpected/disallowed errors/exceptions (default: false)')

        global gConfig
        gConfig = parser.parse_args()
S
Shuduo Sang 已提交
2570

2571
        Logging.clsInit(gConfig)
2572 2573 2574 2575 2576 2577 2578 2579

        Dice.seed(0)  # initial seeding of dice

    def run(self):
        if gConfig.run_tdengine:  # run server
            self.runService()
        else:
            return self.runClient()
S
Steven Li 已提交
2580

2581

2582 2583 2584 2585 2586 2587 2588 2589 2590 2591 2592 2593 2594 2595 2596 2597 2598 2599 2600 2601 2602
class Container():
    _propertyList = {'defTdeInstance'}

    def __init__(self):
        self._cargo = {} # No cargo at the beginning

    def _verifyValidProperty(self, name):
        if not name in self._propertyList:
            raise CrashGenError("Invalid container property: {}".format(name))

    # Called for an attribute, when other mechanisms fail (compare to  __getattribute__)
    def __getattr__(self, name):
        self._verifyValidProperty(name)
        return self._cargo[name] # just a simple lookup

    def __setattr__(self, name, value):
        if name == '_cargo' : # reserved vars
            super().__setattr__(name, value)
            return
        self._verifyValidProperty(name)
        self._cargo[name] = value
S
Shuduo Sang 已提交
2603