Merge pull request #288 from drift-labs/chore/grpc_publisher

remove grpc impl, use sdk grpc OrderSubscriber
This commit is contained in:
wphan
2024-12-03 09:38:50 -08:00
committed by GitHub
3 changed files with 58 additions and 314 deletions

View File

@@ -1,6 +1,5 @@
import { DLOB, OrderSubscriber, UserAccount, UserMap } from '@drift-labs/sdk';
import { PublicKey } from '@solana/web3.js';
import { GeyserOrderSubscriber } from './grpc/OrderSubscriberGRPC';
export type DLOBProvider = {
subscribe(): Promise<void>;
@@ -91,45 +90,3 @@ export function getDLOBProviderFromOrderSubscriber(
},
};
}
export function getDLOBProviderFromGrpcOrderSubscriber(
orderSubscriber: GeyserOrderSubscriber
): DLOBProvider {
return {
subscribe: async () => {
await orderSubscriber.subscribe();
},
getDLOB: async (slot: number) => {
return await orderSubscriber.getDLOB(slot);
},
getUniqueAuthorities: () => {
const authorities = new Set<string>();
for (const { userAccount } of orderSubscriber.usersAccounts.values()) {
authorities.add(userAccount.authority.toBase58());
}
const pubkeys = Array.from(authorities).map((a) => new PublicKey(a));
return pubkeys;
},
getUserAccounts: function* () {
for (const [
key,
{ userAccount },
] of orderSubscriber.usersAccounts.entries()) {
yield { userAccount: userAccount, publicKey: new PublicKey(key) };
}
},
getUserAccount: (publicKey) => {
return orderSubscriber.usersAccounts.get(publicKey.toString())
?.userAccount;
},
size(): number {
return orderSubscriber.usersAccounts.size;
},
fetch() {
return orderSubscriber.fetch();
},
getSlot: () => {
return orderSubscriber.getSlot();
},
};
}

View File

@@ -1,253 +0,0 @@
import { BorshAccountsCoder } from '@coral-xyz/anchor';
import { Commitment, PublicKey, RpcResponseAndContext } from '@solana/web3.js';
import { Buffer } from 'buffer';
import Client, {
CommitmentLevel,
SubscribeRequest,
SubscribeUpdate,
} from '@triton-one/yellowstone-grpc';
import {
BN,
DLOB,
DriftClient,
UserAccount,
getUserFilter,
getUserWithOrderFilter,
} from '@drift-labs/sdk';
import { ClientDuplexStream } from '@grpc/grpc-js';
type grpcDlobConfig = {
endpoint: string;
token: string;
};
export class GeyserOrderSubscriber {
usersAccounts = new Map<string, { slot: number; userAccount: UserAccount }>();
commitment: Commitment;
driftClient: DriftClient;
config: grpcDlobConfig;
stream: ClientDuplexStream<SubscribeRequest, SubscribeUpdate>;
fetchPromise?: Promise<void>;
fetchPromiseResolver: () => void;
mostRecentSlot: number;
constructor(driftClient: DriftClient, config: grpcDlobConfig) {
this.driftClient = driftClient;
this.config = config;
}
public async subscribe(): Promise<void> {
const client = new Client(this.config.endpoint, this.config.token);
this.stream = await client.subscribe();
const request: SubscribeRequest = {
slots: {},
accounts: {
drift: {
owner: [this.driftClient.program.programId.toBase58()],
filters: [
{
memcmp: {
offset: '0',
bytes: BorshAccountsCoder.accountDiscriminator('User'),
},
},
{
memcmp: {
offset: '4350',
bytes: Uint8Array.from([1]),
},
},
],
account: [],
},
},
transactions: {},
blocks: {},
blocksMeta: {},
accountsDataSlice: [],
commitment: CommitmentLevel.PROCESSED,
entry: {},
};
this.stream.on('data', (chunk: any) => {
if (!chunk.account) {
return;
}
const slot = Number(chunk.account.slot);
this.tryUpdateUserAccount(
new PublicKey(chunk.account.account.pubkey).toBase58(),
'grpc',
chunk.account.account.data,
slot
);
});
return new Promise<void>((resolve, reject) => {
this.stream.write(request, (err) => {
if (err === null || err === undefined) {
resolve();
} else {
reject(err);
}
});
}).catch((reason) => {
console.error(reason);
throw reason;
});
}
async fetch(): Promise<void> {
if (this.fetchPromise) {
return this.fetchPromise;
}
this.fetchPromise = new Promise((resolver) => {
this.fetchPromiseResolver = resolver;
});
try {
const rpcRequestArgs = [
this.driftClient.program.programId.toBase58(),
{
commitment: this.commitment,
filters: [getUserFilter(), getUserWithOrderFilter()],
encoding: 'base64',
withContext: true,
},
];
const rpcJSONResponse: any =
// @ts-ignore
await this.driftClient.connection._rpcRequest(
'getProgramAccounts',
rpcRequestArgs
);
const rpcResponseAndContext: RpcResponseAndContext<
Array<{
pubkey: PublicKey;
account: {
data: [string, string];
};
}>
> = rpcJSONResponse.result;
const slot: number = rpcResponseAndContext.context.slot;
const programAccountSet = new Set<string>();
for (const programAccount of rpcResponseAndContext.value) {
const key = programAccount.pubkey.toString();
programAccountSet.add(key);
this.tryUpdateUserAccount(
key,
'raw',
programAccount.account.data,
slot
);
// give event loop a chance to breathe
await new Promise((resolve) => setTimeout(resolve, 0));
}
for (const key of this.usersAccounts.keys()) {
if (!programAccountSet.has(key)) {
this.usersAccounts.delete(key);
}
// give event loop a chance to breathe
await new Promise((resolve) => setTimeout(resolve, 0));
}
} catch (e) {
console.error(e);
} finally {
this.fetchPromiseResolver();
this.fetchPromise = undefined;
}
}
tryUpdateUserAccount(
key: string,
dataType: 'raw' | 'grpc',
data: string[] | Buffer | UserAccount,
slot: number
): void {
if (!this.mostRecentSlot || slot > this.mostRecentSlot) {
this.mostRecentSlot = slot;
}
const slotAndUserAccount = this.usersAccounts.get(key);
if (!slotAndUserAccount || slotAndUserAccount.slot <= slot) {
// Polling leads to a lot of redundant decoding, so we only decode if data is from a fresh slot
let buffer: Buffer;
if (dataType === 'raw') {
buffer = Buffer.from(data[0], data[1]);
} else {
buffer = data as Buffer;
}
const newLastActiveSlot = new BN(
buffer.subarray(4328, 4328 + 8),
undefined,
'le'
);
if (
slotAndUserAccount &&
slotAndUserAccount.userAccount.lastActiveSlot.gt(newLastActiveSlot)
) {
return;
}
const userAccount =
this.driftClient.program.account.user.coder.accounts.decodeUnchecked(
'User',
buffer
) as UserAccount;
if (userAccount.hasOpenOrder) {
this.usersAccounts.set(key, { slot, userAccount });
} else {
this.usersAccounts.delete(key);
}
}
}
public async getDLOB(slot: number): Promise<DLOB> {
const dlob = new DLOB();
for (const [key, { userAccount }] of this.usersAccounts.entries()) {
const userAccountPubkey = new PublicKey(key);
for (const order of userAccount.orders) {
dlob.insertOrder(order, userAccountPubkey.toBase58(), slot);
}
}
return dlob;
}
public getSlot(): number {
return this.mostRecentSlot ?? 0;
}
public async unsubscribe(): Promise<void> {
return new Promise<void>((resolve, reject) => {
this.stream.write(
{
slots: {},
accounts: {},
transactions: {},
blocks: {},
blocksMeta: {},
accountsDataSlice: [],
entry: {},
},
(err) => {
if (err === null || err === undefined) {
resolve();
} else {
reject(err);
}
}
);
}).catch((reason) => {
console.error(reason);
throw reason;
});
}
}

View File

@@ -20,6 +20,7 @@ import {
PhoenixSubscriber,
MarketType,
OraclePriceData,
OrderSubscriberConfig,
} from '@drift-labs/sdk';
import { RedisClient, RedisClientPrefix } from '@drift/common/clients';
@@ -37,12 +38,10 @@ import {
} from '../dlob-subscriber/DLOBSubscriberIO';
import {
DLOBProvider,
getDLOBProviderFromGrpcOrderSubscriber,
getDLOBProviderFromOrderSubscriber,
getDLOBProviderFromUserMap,
} from '../dlobProvider';
import FEATURE_FLAGS from '../utils/featureFlags';
import { GeyserOrderSubscriber } from '../grpc/OrderSubscriberGRPC';
import express, { Response, Request } from 'express';
import { handleHealthCheck } from '../core/metrics';
import { setGlobalDispatcher, Agent } from 'undici';
@@ -76,6 +75,17 @@ const dlobSlotGauge = new Gauge({
'redisClient',
],
});
const oracleSlotGauge = new Gauge({
name: 'oracle_slot',
help: 'Last updated slot of oracle',
labelNames: [
'marketIndex',
'marketType',
'marketName',
'redisPrefix',
'redisClient',
],
});
//@ts-ignore
const sdkConfig = initialize({ env: process.env.ENV });
@@ -116,6 +126,10 @@ const SPOT_MARKETS_TO_LOAD =
logger.info(`RPC endpoint: ${endpoint}`);
logger.info(`WS endpoint: ${wsEndpoint}`);
logger.info(`Token: ${token}`);
logger.info(
`useOrderSubscriber: ${useOrderSubscriber}, useWebsocket: ${useWebsocket}, useGrpc: ${useGrpc}`
);
logger.info(`DriftEnv: ${driftEnv}`);
logger.info(`Commit: ${commitHash}`);
@@ -377,17 +391,38 @@ const main = async () => {
let dlobProvider: DLOBProvider;
if (useOrderSubscriber) {
let subscriptionConfig;
let subscriptionConfig: any = {
type: 'polling',
commitment: stateCommitment,
frequency: ORDERBOOK_UPDATE_INTERVAL,
};
if (useWebsocket) {
subscriptionConfig = {
type: 'websocket',
commitment: stateCommitment,
};
} else {
}
// USE_GRPC=true will override websocket
if (useGrpc) {
if (!token) {
throw new Error('TOKEN is required for grpc');
}
if (!endpoint) {
throw new Error('ENDPOINT is required for grpc');
}
if (useWebsocket) {
logger.warn('USE_GRPC overriding USE_WEBSOCKET');
}
subscriptionConfig = {
type: 'polling',
type: 'grpc',
grpcConfigs: {
endpoint: endpoint,
token: token,
commitmentLevel: stateCommitment,
},
commitment: stateCommitment,
frequency: ORDERBOOK_UPDATE_INTERVAL,
};
}
@@ -401,17 +436,6 @@ const main = async () => {
slotSource = {
getSlot: () => orderSubscriber.getSlot(),
};
} else if (useGrpc) {
const grpcOrderSubscriber = new GeyserOrderSubscriber(driftClient, {
endpoint: endpoint,
token: token,
});
dlobProvider = getDLOBProviderFromGrpcOrderSubscriber(grpcOrderSubscriber);
slotSource = {
getSlot: () => grpcOrderSubscriber.getSlot(),
};
} else {
const userMap = new UserMap({
driftClient,
@@ -463,6 +487,9 @@ const main = async () => {
setInterval(() => {
const slot = slotSource.getSlot();
perpMarketInfos.forEach((market) => {
const oracleDataAndSlot = driftClient.getOracleDataForPerpMarket(
market.marketIndex
);
dlobSlotGauge.set(
{
marketIndex: market.marketIndex,
@@ -473,8 +500,21 @@ const main = async () => {
},
slot
);
oracleSlotGauge.set(
{
marketIndex: market.marketIndex,
marketType: 'perp',
marketName: market.marketName,
redisClient: REDIS_CLIENT,
redisPrefix: RedisClientPrefix[REDIS_CLIENT],
},
oracleDataAndSlot.slot.toNumber()
);
});
spotMarketInfos.forEach((market) => {
const oracleDataAndSlot = driftClient.getOracleDataForSpotMarket(
market.marketIndex
);
dlobSlotGauge.set(
{
marketIndex: market.marketIndex,
@@ -483,7 +523,7 @@ const main = async () => {
redisClient: REDIS_CLIENT,
redisPrefix: RedisClientPrefix[REDIS_CLIENT],
},
slot
oracleDataAndSlot.slot.toNumber()
);
});
}, 10_000);