给mysql增加自定义函数(UDF)，基于雪花算法生成分布式ID

#include <mysql.h> #include <cstdlib> #include <pthread.h> #include <stdio.h> #include <string.h> #include <sys/time.h> #include <unistd.h> #include <vector> #ifndef __SNOWFLAKE__ #define __SNOWFLAKE__ #if defined(SNOWFLAKE_LIBRARY) # define SNOWFLAKE_EXPORT Q_DECL_EXPORT #else # define SNOWFLAKE_EXPORT Q_DECL_IMPORT #endif #ifdef __cplusplus extern "C" { #endif int SNOWFLAKE_init(UDF_INIT *, UDF_ARGS *, char *); void SNOWFLAKE_deinit(UDF_INIT *); long long SNOWFLAKE(UDF_INIT *, UDF_ARGS *); #ifdef __cplusplus } #endif // 起始时间戳，毫秒 #define SNOWFLAKE_EPOCH 1420041600000 // 新时间戳时，最大的随机序列 // 注：序列不包含此数，必须是2的幂次方 #define SNOWFLAKE_MAX_RAND_SEQUENCE 64 // 时钟回拨时，最长等待时间，毫秒 #define SNOWFLAKE_MAX_CLOCK_BACKWARD 10 // 错误码：非法的机器id #define SNOWFLAKE_ERROR_CODE_ILLEGAL_WORKER_ID 1 // 错误码：时钟回拨超时 #define SNOWFLAKE_ERROR_CODE_CLOCK_BACKWARD_TIMEOUT 2 // 错误码：参数太少 #define SNOWFLAKE_ERROR_CODE_TOO_FEW_ARGUMENTS 11 // 错误码：参数太多 #define SNOWFLAKE_ERROR_CODE_TOO_MANY_ARGUMENTS 12 // 错误码：无效的参数类型 #define SNOWFLAKE_ERROR_CODE_INVALID_ARGUMENT_TYPE 13 /** * 共64位 * 0 - 0000000000 0000000000 0000000000 0000000000 0 - 00000 - 00000 - 000000000000 * 正号|1位 - 毫秒|41位 - 数据标识id|5位 - 机器id|5位 - 序列|12位 */ class Snowflake { public: // 错误码：非法的机器id static const int ERROR_CODE_ILLEGAL_WORKER_ID = SNOWFLAKE_ERROR_CODE_ILLEGAL_WORKER_ID; // 错误码：时钟回拨超时 static const int ERROR_CODE_CLOCK_BACKWARD_TIMEOUT = SNOWFLAKE_ERROR_CODE_CLOCK_BACKWARD_TIMEOUT; // 新时间戳时，最大的随机序列 static const int MAX_RAND_SEQUENCE = SNOWFLAKE_MAX_RAND_SEQUENCE; // 新时间戳时，最大的随机序列掩码 static const int MAX_RAND_SEQUENCE_MASK = MAX_RAND_SEQUENCE - 1; // 起始时间戳，毫秒 static const long long EPOCH = SNOWFLAKE_EPOCH; // 时钟回拨时，最长等待时间，毫秒 static const long long MAX_CLOCK_BACKWARD = SNOWFLAKE_MAX_CLOCK_BACKWARD; // 机器id占位 static const int WORKER_ID_BITS = 10; // 序列占位 static const int SEQUENCE_BITS = 12; // 最大机器id static const int MAX_WORKER_ID; // 时间戳，毫秒，左移位数 static const int TIMESTAMP_LEFT_SHIFT; // 机器id，左移位数 static const int WORKER_ID_LEFT_SHIFT; // 序列掩码 static const int SEQUENCE_MASK; private: long long mWorkerId; long long mLastTimestamp; long long mSequence; pthread_mutex_t mMu; public: Snowflake(int workerId); virtual ~Snowflake(); long long nextId(); long long timeGen(); long long sequenceGen(); long long tilNextMillis(long long lastTimestamp); void tilClockForward(long long timestamp, long long lastTimestamp); void sleep(long long msec); long long getWorkerId(); long long getLastTimestamp(); long long getSequence(); }; const int Snowflake::MAX_WORKER_ID = ~(-1 << Snowflake::WORKER_ID_BITS); const int Snowflake::TIMESTAMP_LEFT_SHIFT = Snowflake::SEQUENCE_BITS + Snowflake::WORKER_ID_BITS; const int Snowflake::WORKER_ID_LEFT_SHIFT = Snowflake::SEQUENCE_BITS; const int Snowflake::SEQUENCE_MASK = ~(-1 << Snowflake::SEQUENCE_BITS); class SnowflakeVector { private: std::vector<Snowflake *> mInstances; public: SnowflakeVector(unsigned int size); virtual ~SnowflakeVector(); Snowflake *getInstance(int i); std::vector<Snowflake *> getInstances(); }; class SnowflakeFactory { public: static const int SNOWFLAKE_SIZE; private: static SnowflakeVector sSnowflakeVector; public: static Snowflake *getSnowflake(int workerId); }; const int SnowflakeFactory::SNOWFLAKE_SIZE = Snowflake::MAX_WORKER_ID + 1; SnowflakeVector SnowflakeFactory::sSnowflakeVector(SnowflakeFactory::SNOWFLAKE_SIZE); Snowflake::Snowflake(int workerId) : mWorkerId(workerId), mLastTimestamp(-1), mSequence(0) { if (mWorkerId < 0 || mWorkerId > MAX_WORKER_ID) std::exit(ERROR_CODE_ILLEGAL_WORKER_ID); mMu = PTHREAD_MUTEX_INITIALIZER; } Snowflake::~Snowflake() { } long long Snowflake::nextId() { pthread_mutex_lock(&mMu); long long timestamp = timeGen(); if (timestamp < mLastTimestamp) { tilClockForward(timestamp, mLastTimestamp); timestamp = timeGen(); } if (timestamp == mLastTimestamp) { mSequence = (mSequence + 1) & SEQUENCE_MASK; if (mSequence == 0) timestamp = tilNextMillis(mLastTimestamp); } else { if (MAX_RAND_SEQUENCE > 0) mSequence = sequenceGen(); else mSequence = 0; } mLastTimestamp = timestamp; long long r = ((timestamp - EPOCH) << TIMESTAMP_LEFT_SHIFT) | (mWorkerId << WORKER_ID_LEFT_SHIFT) | mSequence; pthread_mutex_unlock(&mMu); return r; } long long Snowflake::timeGen() { struct timeval tv; gettimeofday(&tv, NULL); return (((long long)tv.tv_sec) * 1000LL) + (tv.tv_usec / 1000L); } long long Snowflake::sequenceGen() { std::srand((unsigned)time(NULL)); int r = std::rand(); return r & MAX_RAND_SEQUENCE_MASK; } long long Snowflake::tilNextMillis(long long lastTimestamp) { long long timestamp = timeGen(); while (timestamp <= lastTimestamp) timestamp = timeGen(); return timestamp; } void Snowflake::tilClockForward(long long timestamp, long long lastTimestamp) { if (timestamp < lastTimestamp) { long long waitDifference = lastTimestamp - timestamp; if (waitDifference > MAX_CLOCK_BACKWARD) std::exit(ERROR_CODE_CLOCK_BACKWARD_TIMEOUT); sleep(waitDifference); } } void Snowflake::sleep(long long msec) { usleep(msec * 1000); } long long Snowflake::getWorkerId() { return mWorkerId; } long long Snowflake::getLastTimestamp() { return mLastTimestamp; } long long Snowflake::getSequence() { return mSequence; } SnowflakeVector::SnowflakeVector(unsigned int size) : mInstances(size, NULL) { for (int i = 0; i < size; i++) mInstances[i] = new Snowflake(i); } SnowflakeVector::~SnowflakeVector() { for (int i = 0; i < mInstances.size(); i++) { if (mInstances[i] != NULL) { delete mInstances[i]; mInstances[i] = NULL; } } mInstances.clear(); } Snowflake *SnowflakeVector::getInstance(int i) { return (i >= 0) ? mInstances[i] : NULL; } std::vector<Snowflake *> SnowflakeVector::getInstances() { return mInstances; } Snowflake *SnowflakeFactory::getSnowflake(int workerId) { return sSnowflakeVector.getInstance(workerId); } int SNOWFLAKE_init(UDF_INIT *initid, UDF_ARGS *args, char *message) { if (args->arg_count == 0) { strcpy(message, "too few arguments: expected 1"); return SNOWFLAKE_ERROR_CODE_TOO_FEW_ARGUMENTS; } if (args->arg_count > 1) { sprintf(message, "too many arguments: expected 1, have %d", args->arg_count); return SNOWFLAKE_ERROR_CODE_TOO_MANY_ARGUMENTS; } if (args->arg_type[0] != INT_RESULT) { strcpy(message, "invalid argument type: expected INT"); return SNOWFLAKE_ERROR_CODE_INVALID_ARGUMENT_TYPE; } const long long input = *((long long *)args->args[0]); if (input < 0 || input > Snowflake::MAX_WORKER_ID) { sprintf(message, "illegal worker id: expected range