From 9862aa8bed51640a9028e7216e86abf52fbc23e2 Mon Sep 17 00:00:00 2001
From: "A.J. Beamon" <ajbeamon@apple.com>
Date: Wed, 4 Mar 2020 11:15:32 -0800
Subject: [PATCH] Add support for setting knobs in fdbcli

---
 fdbcli/fdbcli.actor.cpp | 50 ++++++++++++++++++++++++++++++++++++++---
 1 file changed, 47 insertions(+), 3 deletions(-)

diff --git a/fdbcli/fdbcli.actor.cpp b/fdbcli/fdbcli.actor.cpp
index cf76fe7ee4..0697fe821d 100644
--- a/fdbcli/fdbcli.actor.cpp
+++ b/fdbcli/fdbcli.actor.cpp
@@ -69,7 +69,8 @@ enum {
 	OPT_NO_STATUS,
 	OPT_STATUS_FROM_JSON,
 	OPT_VERSION,
-	OPT_TRACE_FORMAT
+	OPT_TRACE_FORMAT,
+	OPT_KNOB
 };
 
 CSimpleOpt::SOption g_rgOptions[] = { { OPT_CONNFILE, "-C", SO_REQ_SEP },
@@ -87,12 +88,13 @@ CSimpleOpt::SOption g_rgOptions[] = { { OPT_CONNFILE, "-C", SO_REQ_SEP },
 	                                  { OPT_VERSION, "--version", SO_NONE },
 	                                  { OPT_VERSION, "-v", SO_NONE },
 	                                  { OPT_TRACE_FORMAT, "--trace_format", SO_REQ_SEP },
+	                                  { OPT_KNOB, "--knob_", SO_REQ_SEP },
 
 #ifndef TLS_DISABLED
 	                                  TLS_OPTION_FLAGS
 #endif
 
-	                                      SO_END_OF_OPTIONS };
+	                                  SO_END_OF_OPTIONS };
 
 void printAtCol(const char* text, int col) {
 	const char* iter = text;
@@ -423,6 +425,8 @@ static void printProgramUsage(const char* name) {
 #ifndef TLS_DISABLED
 	       TLS_HELP
 #endif
+	       "  --knob_KNOBNAME KNOBVALUE\n"
+	       "                 Changes a knob option. KNOBNAME should be lowercase.\n"
 	       "  -v, --version  Print FoundationDB CLI version information and exit.\n"
 	       "  -h, --help     Display this help and exit.\n");
 }
@@ -2444,6 +2448,8 @@ struct CLIOptions {
 	std::string tlsCAPath;
 	std::string tlsPassword;
 
+	std::vector<std::pair<std::string, std::string>> knobs;
+
 	CLIOptions( int argc, char* argv[] )
 		: trace(false),
 		 exit_timeout(0),
@@ -2467,9 +2473,37 @@ struct CLIOptions {
 		}
 		if (exit_timeout && !exec.present()) {
 			fprintf(stderr, "ERROR: --timeout may only be specified with --exec\n");
-			exit_code = 1;
+			exit_code = FDB_EXIT_ERROR;
 			return;
 		}
+
+		delete FLOW_KNOBS;
+		FlowKnobs* flowKnobs = new FlowKnobs(true);
+		FLOW_KNOBS = flowKnobs;
+
+		delete CLIENT_KNOBS;
+		ClientKnobs* clientKnobs = new ClientKnobs(true);
+		CLIENT_KNOBS = clientKnobs;
+
+		for(auto k=knobs.begin(); k!=knobs.end(); ++k) {
+			try {
+				if (!flowKnobs->setKnob( k->first, k->second ) &&
+					!clientKnobs->setKnob( k->first, k->second ))
+				{
+					fprintf(stderr, "ERROR: Unrecognized knob option '%s'\n", k->first.c_str());
+					exit_code = FDB_EXIT_ERROR;
+				}
+			} catch (Error& e) {
+				if (e.code() == error_code_invalid_option_value) {
+					fprintf(stderr, "ERROR: Invalid value '%s' for knob option '%s'\n", k->second.c_str(), k->first.c_str());
+					exit_code = FDB_EXIT_ERROR;
+				}
+				else {
+					fprintf(stderr, "ERROR: Failed to set knob option '%s': %s\n", k->first.c_str(), e.what());
+					exit_code = FDB_EXIT_ERROR;
+				}
+			}
+		}
 	}
 
 	int processArg(CSimpleOpt& args) {
@@ -2536,6 +2570,16 @@ struct CLIOptions {
 			    }
 			    traceFormat = args.OptionArg();
 			    break;
+			case OPT_KNOB: {
+				std::string syn = args.OptionSyntax();
+				if (!StringRef(syn).startsWith(LiteralStringRef("--knob_"))) {
+					fprintf(stderr, "ERROR: unable to parse knob option '%s'\n", syn.c_str());
+					return FDB_EXIT_ERROR;
+				}
+				syn = syn.substr(7);
+				knobs.push_back( std::make_pair( syn, args.OptionArg() ) );
+				break;
+			}
 		    case OPT_VERSION:
 			    printVersion();
 			    return FDB_EXIT_SUCCESS;