Improvements after review

src/app/app.module.ts

@@ -46,9 +46,9 @@ import { StakingRouterModule } from '../staking-router-modules';
     }),
     ScheduleModule.forRoot(),
     KeyRegistryModule.forRootAsync({
-      inject: [SimpleFallbackJsonRpcBatchProvider],
-      async useFactory(provider) {
-        return { provider };
+      inject: [SimpleFallbackJsonRpcBatchProvider, ConfigService],
+      async useFactory(provider, configService) {
+        return { provider, keysBatchSize: configService?.get('KEYS_FETCH_BATCH_SIZE') };
       },
     }),
     StakingRouterModule,

src/common/config/env.validation.ts

@@ -11,6 +11,7 @@ import {
   IsBoolean,
   ValidateIf,
   IsNotEmpty,
+  IsPositive,
 } from 'class-validator';
 import { Environment, LogLevel, LogFormat } from './interfaces';
 import { NonEmptyArray } from '@lido-nestjs/execution/dist/interfaces/non-empty-array';
@@ -169,6 +170,11 @@ export class EnvironmentVariables {
   @IsOptional()
   @IsString()
   LIDO_LOCATOR_ADDRESS = '';
+
+  @IsOptional()
+  @IsPositive()
+  @Transform(({ value }) => parseInt(value, 10))
+  KEYS_FETCH_BATCH_SIZE = 1100;
 }
 
 export function validate(config: Record<string, unknown>) {

src/common/registry/fetch/interfaces/module.interface.ts

@@ -3,10 +3,13 @@ import { ModuleMetadata } from '@nestjs/common';
 import { Signer } from 'ethers';
 import { Provider } from '@ethersproject/providers';
 
+export const REGISTRY_FETCH_OPTIONS_TOKEN = Symbol('registryFetchOptionsToken');
+
 export interface RegistryFetchOptions {
   registryAddress?: string;
   lidoAddress?: string;
   provider?: Provider | Signer;
+  keysBatchSize?: number;
 }
 
 export interface RegistryFetchModuleSyncOptions extends Pick<ModuleMetadata, 'imports'>, RegistryFetchOptions {}

src/common/registry/fetch/key-batch.constants.ts

@@ -5,3 +5,4 @@
  */
 export const KEYS_LENGTH = 96;
 export const SIGNATURE_LENGTH = 192;
+export const KEYS_BATCH_SIZE = 1100;

src/common/registry/fetch/key-batch.fetch.ts

@@ -3,13 +3,15 @@ import { REGISTRY_CONTRACT_TOKEN, Registry } from '@lido-nestjs/contracts';
 import { CallOverrides } from './interfaces/overrides.interface';
 import { KeyBatchRecord, RegistryKey } from './interfaces/key.interface';
 import { RegistryOperatorFetchService } from './operator.fetch';
-import { KEYS_LENGTH, SIGNATURE_LENGTH } from './key-batch.constants';
+import { KEYS_BATCH_SIZE, KEYS_LENGTH, SIGNATURE_LENGTH } from './key-batch.constants';
+import { RegistryFetchOptions, REGISTRY_FETCH_OPTIONS_TOKEN } from './interfaces/module.interface';
 
 @Injectable()
 export class RegistryKeyBatchFetchService {
   constructor(
     protected readonly operatorsService: RegistryOperatorFetchService,
     @Inject(REGISTRY_CONTRACT_TOKEN) private contract: Registry,
+    @Inject(REGISTRY_FETCH_OPTIONS_TOKEN) private options: RegistryFetchOptions,
   ) {}
 
   private getContract(moduleAddress: string) {
@@ -107,8 +109,7 @@ export class RegistryKeyBatchFetchService {
     fromIndex: number,
     totalAmount: number,
   ) {
-    // TODO: move to constants/config cause this limit depends on eth node
-    const batchSize = 1100;
+    const batchSize = this.options.keysBatchSize || KEYS_BATCH_SIZE;
 
     const numberOfBatches = Math.ceil(totalAmount / batchSize);
     const promises: Promise<RegistryKey[]>[] = [];

src/common/registry/fetch/meta.fetch.ts

@@ -1,6 +1,5 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { Inject, Injectable } from '@nestjs/common';
-// import { Registry__factory } from '@lido-nestjs/contracts';
 import { CallOverrides } from './interfaces/overrides.interface';
 import { REGISTRY_CONTRACT_TOKEN, Registry } from '@lido-nestjs/contracts';
 
@@ -12,10 +11,9 @@ export class RegistryMetaFetchService {
     return this.contract.attach(moduleAddress);
   }
 
-  /** fetches keys operation index */
-  public async fetchKeysOpIndex(moduleAddress: string, overrides: CallOverrides = {}): Promise<number> {
-    // TODO: read data from all contract that implement curated-v1-onchain type
-    const bigNumber = await this.getContract(moduleAddress).getKeysOpIndex(overrides as any);
+  /** Fetches nonce from staking module contract */
+  public async fetchStakingModuleNonce(moduleAddress: string, overrides: CallOverrides = {}): Promise<number> {
+    const bigNumber = await this.getContract(moduleAddress).getNonce(overrides as any);
     return bigNumber.toNumber();
   }
 }

src/common/registry/fetch/registry-fetch.module.ts

@@ -1,6 +1,10 @@
 import { DynamicModule, Module } from '@nestjs/common';
 import { LidoContractModule, RegistryContractModule } from '@lido-nestjs/contracts';
-import { RegistryFetchModuleSyncOptions, RegistryFetchModuleAsyncOptions } from './interfaces/module.interface';
+import {
+  RegistryFetchModuleSyncOptions,
+  RegistryFetchModuleAsyncOptions,
+  REGISTRY_FETCH_OPTIONS_TOKEN,
+} from './interfaces/module.interface';
 import { RegistryOperatorFetchService } from './operator.fetch';
 import { RegistryMetaFetchService } from './meta.fetch';
 import { RegistryKeyFetchService } from './key.fetch';
@@ -52,6 +56,12 @@ export class RegistryFetchModule {
           provider: options?.provider,
         }),
       ],
+      providers: [
+        {
+          provide: REGISTRY_FETCH_OPTIONS_TOKEN,
+          useValue: options?.keysBatchSize ? { keysBatchSize: options.keysBatchSize } : {},
+        },
+      ],
       exports: [LidoContractModule, RegistryContractModule],
     };
   }
@@ -80,6 +90,13 @@ export class RegistryFetchModule {
           inject: options.inject,
         }),
       ],
+      providers: [
+        {
+          provide: REGISTRY_FETCH_OPTIONS_TOKEN,
+          useFactory: options.useFactory,
+          inject: options.inject,
+        },
+      ],
       exports: [LidoContractModule, RegistryContractModule],
     };
   }

src/common/registry/main/abstract-registry.ts

@@ -110,7 +110,6 @@ export abstract class AbstractRegistryService {
       // TODO: use feature flag
       const result = await this.keyBatchFetch.fetch(moduleAddress, operatorIndex, fromIndex, toIndex, overrides);
 
-      // add moduleAddress
       const operatorKeys = result.filter((key) => key);
 
       this.logger.log('Keys fetched', {
@@ -133,12 +132,10 @@ export abstract class AbstractRegistryService {
 
   /** returns the latest operators data from the db */
   public async getOperatorsFromStorage(moduleAddress: string) {
-    // TODO: find for module
     return await this.operatorStorage.findAll(moduleAddress);
   }
 
   /** returns all the keys from storage */
-  // the same
   public async getOperatorsKeysFromStorage(moduleAddress: string) {
     return await this.keyStorage.findAll(moduleAddress);
   }
@@ -161,12 +158,12 @@ export abstract class AbstractRegistryService {
 
   /** contract */
   /** returns the meta data from the contract */
-  public async getNonceFromContract(moduleAddress: string, blockHash: string): Promise<number> {
-    const keysOpIndex = await this.metaFetch.fetchKeysOpIndex(moduleAddress, { blockTag: { blockHash } });
+  public async getStakingModuleNonce(moduleAddress: string, blockHash: string): Promise<number> {
+    const keysOpIndex = await this.metaFetch.fetchStakingModuleNonce(moduleAddress, { blockTag: { blockHash } });
     return keysOpIndex;
   }
 
-  /** saves all the data to the db */
+  /** saves all data to the db for staking module*/
   public async saveOperators(moduleAddress: string, currentOperators: RegistryOperator[]) {
     // save all data in a transaction
     await this.entityManager.transactional(async (entityManager) => {
@@ -189,20 +186,11 @@ export abstract class AbstractRegistryService {
           await entityManager
             .createQueryBuilder(RegistryOperator)
             .insert(operatorsChunk)
-            // TODO: module_address or moduleAddress ?
             .onConflict(['index', 'module_address'])
             .merge()
             .execute();
         }),
       );
     });
   }
-
-  /** clears the db */
-  public async clear() {
-    await this.entityManager.transactional(async (entityManager) => {
-      entityManager.nativeDelete(RegistryKey, {});
-      entityManager.nativeDelete(RegistryOperator, {});
-    });
-  }
 }

src/common/registry/main/key-registry/key-registry.service.ts

@@ -7,7 +7,7 @@ export class KeyRegistryService extends AbstractRegistryService {
     return currOperator.totalSigningKeys;
   }
   /** returns all operators keys from the db */
-  public async getAllKeysFromStorage(moduleAddress: string) {
+  public async getModuleKeysFromStorage(moduleAddress: string) {
     return await this.keyStorage.findAll(moduleAddress);
   }
   /** returns used keys from the db */

src/common/registry/storage/constants.ts

@@ -1,4 +1,3 @@
 export const KEY_LEN = 98;
 export const DEPOSIT_SIGNATURE_LEN = 194;
-export const MODULE_ADDRESS_LEN = 42;
-export const REWARD_ADDRESS_LEN = 42;
+export const ADDRESS_LEN = 42;

src/common/registry/storage/key.entity.ts

@@ -1,5 +1,5 @@
 import { Entity, EntityRepositoryType, PrimaryKey, PrimaryKeyType, Property } from '@mikro-orm/core';
-import { DEPOSIT_SIGNATURE_LEN, MODULE_ADDRESS_LEN, KEY_LEN } from './constants';
+import { DEPOSIT_SIGNATURE_LEN, ADDRESS_LEN, KEY_LEN } from './constants';
 import { RegistryKeyRepository } from './key.repository';
 
 @Entity({ customRepository: () => RegistryKeyRepository })
@@ -32,6 +32,6 @@ export class RegistryKey {
   used!: boolean;
 
   @PrimaryKey()
-  @Property({ length: MODULE_ADDRESS_LEN })
+  @Property({ length: ADDRESS_LEN })
   moduleAddress!: string;
 }

src/common/registry/storage/operator.entity.ts

@@ -1,5 +1,5 @@
 import { Entity, EntityRepositoryType, PrimaryKey, PrimaryKeyType, Property } from '@mikro-orm/core';
-import { MODULE_ADDRESS_LEN, REWARD_ADDRESS_LEN } from './constants';
+import { ADDRESS_LEN } from './constants';
 import { RegistryOperatorRepository } from './operator.repository';
 
 @Entity({ customRepository: () => RegistryOperatorRepository })
@@ -28,7 +28,7 @@ export class RegistryOperator {
   @Property({ length: 256 })
   name!: string;
 
-  @Property({ length: REWARD_ADDRESS_LEN })
+  @Property({ length: ADDRESS_LEN })
   rewardAddress!: string;
 
   @Property()
@@ -44,6 +44,6 @@ export class RegistryOperator {
   usedSigningKeys!: number;
 
   @PrimaryKey()
-  @Property({ length: MODULE_ADDRESS_LEN })
+  @Property({ length: ADDRESS_LEN })
   moduleAddress!: string;
 }

src/common/registry/test/fetch/meta.fetch.e2e-spec.ts

@@ -23,7 +23,7 @@ describe('Operators', () => {
   });
 
   test('fetch keysOpIndex', async () => {
-    const keysOpIndex = await fetchService.fetchKeysOpIndex(address);
+    const keysOpIndex = await fetchService.fetchStakingModuleNonce(address);
     expect(typeof keysOpIndex).toBe('number');
     expect(keysOpIndex).toBeGreaterThan(0);
   });

src/common/registry/test/fetch/meta.fetch.spec.ts

@@ -36,7 +36,7 @@ describe('Meta', () => {
         const iface = new Interface(Registry__factory.abi);
         return iface.encodeFunctionResult('getKeysOpIndex', [expected]);
       });
-      const result = await fetchService.fetchKeysOpIndex(address);
+      const result = await fetchService.fetchStakingModuleNonce(address);
 
       expect(result).toEqual(expected);
       expect(mockCall).toBeCalledTimes(1);

src/common/registry/test/fetch/operator.fetch.spec.ts

@@ -43,7 +43,6 @@ describe('Operators', () => {
 
     mockCall.mockImplementation(async () => {
       const iface = new Interface(Registry__factory.abi);
-      // TODO: moduleAddress depends on chain id
       operator['moduleAddress'] = address;
       return iface.encodeFunctionResult('getNodeOperator', operatorFields(operator));
     });

src/common/registry/test/key-registry/async.spec.ts

@@ -6,6 +6,7 @@ import { getNetwork } from '@ethersproject/networks';
 import { getDefaultProvider } from '@ethersproject/providers';
 import { KeyRegistryModule, KeyRegistryService, RegistryStorageService } from '../../';
 import { MikroORM } from '@mikro-orm/core';
+import { mikroORMConfig } from '../testing.utils';
 
 describe('Async module initializing', () => {
   const provider = getDefaultProvider('mainnet');
@@ -26,12 +27,7 @@ describe('Async module initializing', () => {
 
   test('forRootAsync', async () => {
     await testModules([
-      MikroOrmModule.forRoot({
-        dbName: ':memory:',
-        type: 'sqlite',
-        allowGlobalContext: true,
-        entities: ['./**/*.entity.ts'],
-      }),
+      MikroOrmModule.forRoot(mikroORMConfig),
       LoggerModule.forRoot({ transports: [nullTransport()] }),
       KeyRegistryModule.forRootAsync({
         async useFactory() {

src/common/registry/test/key-registry/connect.e2e-spec.ts

@@ -5,7 +5,7 @@ import { BatchProviderModule, ExtendedJsonRpcBatchProvider } from '@lido-nestjs/
 
 import { KeyRegistryModule, KeyRegistryService, RegistryStorageService } from '../../';
 
-import { compareTestMetaOperators } from '../testing.utils';
+import { clearDb, compareTestOperators, mikroORMConfig } from '../testing.utils';
 
 import { operators } from '../fixtures/connect.fixture';
 import { MikroORM } from '@mikro-orm/core';
@@ -17,6 +17,7 @@ dotenv.config();
 describe('Registry', () => {
   let registryService: KeyRegistryService;
   let storageService: RegistryStorageService;
+  let mikroOrm: MikroORM;
   if (!process.env.CHAIN_ID) {
     console.error("CHAIN_ID wasn't provides");
     process.exit(1);
@@ -28,12 +29,7 @@ describe('Registry', () => {
 
   beforeEach(async () => {
     const imports = [
-      MikroOrmModule.forRoot({
-        dbName: ':memory:',
-        type: 'sqlite',
-        allowGlobalContext: true,
-        entities: ['./**/*.entity.ts'],
-      }),
+      MikroOrmModule.forRoot(mikroORMConfig),
       BatchProviderModule.forRoot({
         url: process.env.PROVIDERS_URLS as string,
         requestPolicy: {
@@ -53,13 +49,13 @@ describe('Registry', () => {
     const moduleRef = await Test.createTestingModule({ imports }).compile();
     registryService = moduleRef.get(KeyRegistryService);
     storageService = moduleRef.get(RegistryStorageService);
-
-    const generator = moduleRef.get(MikroORM).getSchemaGenerator();
+    mikroOrm = moduleRef.get(MikroORM);
+    const generator = mikroOrm.getSchemaGenerator();
     await generator.updateSchema();
   });
 
   afterEach(async () => {
-    await registryService.clear();
+    await clearDb(mikroOrm);
     await storageService.onModuleDestroy();
   });
 
@@ -68,7 +64,7 @@ describe('Registry', () => {
 
     await registryService.update(address, blockHash);
 
-    await compareTestMetaOperators(address, registryService, {
+    await compareTestOperators(address, registryService, {
       operators: operatorsWithModuleAddress,
     });