Skip to content

Commit

Permalink
Supported static storage
Browse files Browse the repository at this point in the history
  • Loading branch information
GrigoriyPA committed Jan 6, 2025
1 parent 7f2a682 commit a16614d
Show file tree
Hide file tree
Showing 11 changed files with 93 additions and 15 deletions.
2 changes: 2 additions & 0 deletions ydb/core/testlib/basics/helpers.h
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@ namespace NFake {
ui64 SectorSize = 0;
ui64 ChunkSize = 0;
ui64 DiskSize = 0;
bool FormatDisk = true;
TString DiskPath;
};

struct INode {
Expand Down
4 changes: 2 additions & 2 deletions ydb/core/testlib/basics/storage.h
Original file line number Diff line number Diff line change
Expand Up @@ -54,15 +54,15 @@ namespace NKikimr {

static ui64 keySalt = 0;
ui64 salt = ++keySalt;
TString baseDir = Runtime.GetTempDir();
TString baseDir = conf.DiskPath ? conf.DiskPath : Runtime.GetTempDir();

if (Conf.UseDisk) {
MakeDirIfNotExist(baseDir.c_str());
}

PDiskPath = TStringBuilder() << baseDir << "pdisk_1.dat";

if (!Mock) {
if (!Mock && conf.FormatDisk) {
FormatPDisk(PDiskPath,
Conf.DiskSize, Conf.SectorSize, Conf.ChunkSize, PDiskGuid,
0x123 + salt, 0x456 + salt, 0x789 + salt, mainKey,
Expand Down
9 changes: 7 additions & 2 deletions ydb/core/testlib/test_client.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -614,18 +614,21 @@ namespace Tests {

NKikimrBlobStorage::TDefineBox boxConfig;
boxConfig.SetBoxId(Settings->BOX_ID);
boxConfig.SetItemConfigGeneration(Settings->StorageGeneration);

ui32 nodeId = Runtime->GetNodeId(0);
Y_ABORT_UNLESS(nodesInfo->Nodes[0].NodeId == nodeId);
auto& nodeInfo = nodesInfo->Nodes[0];

NKikimrBlobStorage::TDefineHostConfig hostConfig;
hostConfig.SetHostConfigId(nodeId);
hostConfig.SetItemConfigGeneration(Settings->StorageGeneration);
TString path;
if (Settings->UseSectorMap) {
path ="SectorMap:test-client[:2000]";
} else {
path = TStringBuilder() << Runtime->GetTempDir() << "pdisk_1.dat";
TString diskPath = Settings->CustomDiskParams.DiskPath;
path = TStringBuilder() << (diskPath ? diskPath : Runtime->GetTempDir()) << "pdisk_1.dat";
}
hostConfig.AddDrive()->SetPath(path);
if (Settings->Verbose) {
Expand All @@ -641,7 +644,9 @@ namespace Tests {

for (const auto& [poolKind, storagePool] : Settings->StoragePoolTypes) {
if (storagePool.GetNumGroups() > 0) {
bsConfigureRequest->Record.MutableRequest()->AddCommand()->MutableDefineStoragePool()->CopyFrom(storagePool);
auto* command = bsConfigureRequest->Record.MutableRequest()->AddCommand()->MutableDefineStoragePool();
command->CopyFrom(storagePool);
command->SetItemConfigGeneration(Settings->StorageGeneration);
}
}

Expand Down
2 changes: 2 additions & 0 deletions ydb/core/testlib/test_client.h
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,7 @@ namespace Tests {
TString DomainName = TestDomainName;
ui32 NodeCount = 1;
ui32 DynamicNodeCount = 0;
ui64 StorageGeneration = 0;
NFake::TStorage CustomDiskParams;
TControls Controls;
TAppPrepare::TFnReg FrFactory = &DefaultFrFactory;
Expand Down Expand Up @@ -179,6 +180,7 @@ namespace Tests {
TServerSettings& SetDomainName(const TString& value);
TServerSettings& SetNodeCount(ui32 value) { NodeCount = value; return *this; }
TServerSettings& SetDynamicNodeCount(ui32 value) { DynamicNodeCount = value; return *this; }
TServerSettings& SetStorageGeneration(ui64 value) { StorageGeneration = value; return *this; }
TServerSettings& SetCustomDiskParams(const NFake::TStorage& value) { CustomDiskParams = value; return *this; }
TServerSettings& SetControls(const TControls& value) { Controls = value; return *this; }
TServerSettings& SetFrFactory(const TAppPrepare::TFnReg& value) { FrFactory = value; return *this; }
Expand Down
2 changes: 2 additions & 0 deletions ydb/tests/tools/kqprun/.gitignore
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
storage
sync_dir
example
udfs

*.log
*.json
*.sql
Expand Down
26 changes: 20 additions & 6 deletions ydb/tests/tools/kqprun/kqprun.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,7 @@ struct TExecutionOptions {
ValidateScriptExecutionOptions(runnerOptions);
ValidateAsyncOptions(runnerOptions.YdbSettings.AsyncQueriesSettings);
ValidateTraceOpt(runnerOptions);
ValidateStorageSettings(runnerOptions.YdbSettings);
}

private:
Expand Down Expand Up @@ -258,6 +259,19 @@ struct TExecutionOptions {
}
}

static void ValidateStorageSettings(const TYdbSetupSettings& ydbSettings) {
if (ydbSettings.DisableDiskMock) {
if (ydbSettings.NodeCount + ydbSettings.SharedTenants.size() + ydbSettings.DedicatedTenants.size() > 1) {
ythrow yexception() << "Disable disk mock cannot be used for multi node clusters (already disabled)";
} else if (ydbSettings.PDisksPath) {
ythrow yexception() << "Disable disk mock cannot be used with real PDisks (already disabled)";
}
}
if (ydbSettings.FormatStorage && !ydbSettings.PDisksPath) {
ythrow yexception() << "Cannot format storage without real PDisks, please use --storage-path";
}
}

private:
static void ReplaceYqlTokenTemplate(TString& sql) {
const TString variableName = TStringBuilder() << "${" << YQL_TOKEN_VARIABLE << "}";
Expand Down Expand Up @@ -857,9 +871,13 @@ class TMain : public TMainClassArgs {
return static_cast<ui64>(diskSize) << 30;
});

options.AddLongOption("real-pdisks", "Use real PDisks instead of in memory PDisks (also disable disk mock)")
options.AddLongOption("storage-path", "Use real PDisks by specified path instead of in memory PDisks (also disable disk mock), use '-' to use temp directory")
.RequiredArgument("directory")
.StoreResult(&RunnerOptions.YdbSettings.PDisksPath);

options.AddLongOption("format-storage", "Clear storage if it exists on --storage-path")
.NoArgument()
.SetFlag(&RunnerOptions.YdbSettings.UseRealPDisks);
.SetFlag(&RunnerOptions.YdbSettings.FormatStorage);

options.AddLongOption("disable-disk-mock", "Disable disk mock on single node cluster")
.NoArgument()
Expand All @@ -882,10 +900,6 @@ class TMain : public TMainClassArgs {
int DoRun(NLastGetopt::TOptsParseResult&&) override {
ExecutionOptions.Validate(RunnerOptions);

if (RunnerOptions.YdbSettings.DisableDiskMock && RunnerOptions.YdbSettings.NodeCount + RunnerOptions.YdbSettings.SharedTenants.size() + RunnerOptions.YdbSettings.DedicatedTenants.size() > 1) {
ythrow yexception() << "Disable disk mock cannot be used for multi node clusters";
}

RunnerOptions.YdbSettings.YqlToken = YqlToken;
RunnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry(UdfsDirectory, UdfsPaths, ExcludeLinkedUdfs).Get();

Expand Down
3 changes: 2 additions & 1 deletion ydb/tests/tools/kqprun/src/common.h
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,8 @@ struct TYdbSetupSettings {
bool SameSession = false;

bool DisableDiskMock = false;
bool UseRealPDisks = false;
bool FormatStorage = false;
std::optional<TString> PDisksPath;
ui64 DiskSize = 32_GB;

bool MonitoringEnabled = false;
Expand Down
7 changes: 7 additions & 0 deletions ydb/tests/tools/kqprun/src/proto/storage_meta.proto
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
syntax = "proto3";

package NKqpRun;

message TStorageMeta {
uint64 StorageGeneration = 1;
}
9 changes: 9 additions & 0 deletions ydb/tests/tools/kqprun/src/proto/ya.make
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
PROTO_LIBRARY()

ONLY_TAGS(CPP_PROTO)

SRCS(
storage_meta.proto
)

END()
2 changes: 2 additions & 0 deletions ydb/tests/tools/kqprun/src/ya.make
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ SRCS(

PEERDIR(
ydb/core/testlib

ydb/tests/tools/kqprun/src/proto
)

GENERATE_ENUM_SERIALIZATION(common.h)
Expand Down
42 changes: 38 additions & 4 deletions ydb/tests/tools/kqprun/src/ydb_setup.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,9 @@
#include <ydb/core/testlib/test_client.h>

#include <ydb/library/yql/providers/s3/actors/yql_s3_actors_factory_impl.h>

#include <ydb/tests/tools/kqprun/src/proto/storage_meta.pb.h>

#include <yql/essentials/utils/log/log.h>


Expand Down Expand Up @@ -173,15 +176,46 @@ class TYdbSetup::TImpl {
}

void SetStorageSettings(NKikimr::Tests::TServerSettings& serverSettings) const {
TString diskPath;
if (Settings_.PDisksPath && *Settings_.PDisksPath != "-") {
diskPath = TStringBuilder() << *Settings_.PDisksPath << "/";
}

bool formatDisk = true;
NKqpRun::TStorageMeta storageMeta;
if (diskPath) {
TFsPath storageMetaPath(diskPath);
storageMetaPath.MkDirs();

storageMetaPath = storageMetaPath.Child("kqprun_storage_meta.conf");
if (storageMetaPath.Exists() && !Settings_.FormatStorage) {
if (!google::protobuf::TextFormat::ParseFromString(TFileInput(storageMetaPath.GetPath()).ReadAll(), &storageMeta)) {
ythrow yexception() << "Storage meta is corrupted, please use --format-storage";
}
storageMeta.SetStorageGeneration(storageMeta.GetStorageGeneration() + 1);
formatDisk = false;
}

TString storageMetaStr;
google::protobuf::TextFormat::PrintToString(storageMeta, &storageMetaStr);

TFileOutput storageMetaOutput(storageMetaPath.GetPath());
storageMetaOutput.Write(storageMetaStr);
storageMetaOutput.Finish();
}

const NKikimr::NFake::TStorage storage = {
.UseDisk = Settings_.UseRealPDisks,
.UseDisk = !!Settings_.PDisksPath,
.SectorSize = NKikimr::TTestStorageFactory::SECTOR_SIZE,
.ChunkSize = Settings_.UseRealPDisks ? NKikimr::TTestStorageFactory::CHUNK_SIZE : NKikimr::TTestStorageFactory::MEM_CHUNK_SIZE,
.DiskSize = Settings_.DiskSize
.ChunkSize = Settings_.PDisksPath ? NKikimr::TTestStorageFactory::CHUNK_SIZE : NKikimr::TTestStorageFactory::MEM_CHUNK_SIZE,
.DiskSize = Settings_.DiskSize,
.FormatDisk = formatDisk,
.DiskPath = diskPath
};

serverSettings.SetEnableMockOnSingleNode(!Settings_.DisableDiskMock && !Settings_.UseRealPDisks);
serverSettings.SetEnableMockOnSingleNode(!Settings_.DisableDiskMock && !Settings_.PDisksPath);
serverSettings.SetCustomDiskParams(storage);
serverSettings.SetStorageGeneration(storageMeta.GetStorageGeneration());
}

NKikimr::Tests::TServerSettings GetServerSettings(ui32 grpcPort) {
Expand Down

0 comments on commit a16614d

Please sign in to comment.