Merge e5ef77f into d42d275

src/plugins/mf_priority.cpp

@@ -350,6 +350,34 @@ static bool check_map_for_dne_only ()
 }
 
 
+/*
+ * Update the jobspec with the default bank the association used to
+ * submit their job under.
+ */
+static int update_jobspec_bank (flux_plugin_t *p, int userid)
+{
+    char *bank = NULL;
+    std::map<int, std::map<std::string, struct bank_info>>::iterator it;
+
+    it = users.find (userid);
+    if (it == users.end ()) {
+        return -1;
+    }
+
+    // look up default bank
+    bank = const_cast<char*> (users_def_bank[userid].c_str ());
+
+    // post jobspec-update event
+    if (flux_jobtap_jobspec_update_pack (p,
+                                         "{s:s}",
+                                         "attributes.system.bank",
+                                         bank) < 0)
+        return -1;
+
+    return 0;
+}
+
+
 /******************************************************************************
  *                                                                            *
  *                               Callbacks                                    *
@@ -639,6 +667,17 @@ static int priority_cb (flux_plugin_t *p,
                                          &bank_it->second,
                                          NULL) < 0)
                 flux_log_error (h, "flux_jobtap_job_aux_set");
+
+            // now that we know the user/bank info associated with this job,
+            // we need to update jobspec with the default bank used to
+            // submit this job under
+            if (update_jobspec_bank (p, userid) < 0) {
+                flux_jobtap_raise_exception (p, FLUX_JOBTAP_CURRENT_JOB,
+                                             "mf_priority", 0,
+                                             "failed to update jobspec "
+                                             "with bank name");
+                return -1;
+            }
         }
     }
 
@@ -853,6 +892,14 @@ static int new_cb (flux_plugin_t *p,
                                              "entry does not exist");
                 return -1;
             }
+            // update jobspec with default bank
+            if (update_jobspec_bank (p, userid) < 0) {
+                flux_jobtap_raise_exception (p, FLUX_JOBTAP_CURRENT_JOB,
+                                            "mf_priority", 0,
+                                            "failed to update jobspec "
+                                            "with bank name");
+                return -1;
+            }
         }
 
         max_run_jobs = bank_it->second.max_run_jobs;

t/Makefile.am

@@ -30,6 +30,7 @@ TESTSCRIPTS = \
 	t1026-flux-account-perms.t \
 	t1027-mf-priority-issue376.t \
 	t1028-mf-priority-issue385.t \
+	t1029-mf-priority-default-bank.t \
 	t5000-valgrind.t \
 	python/t1000-example.py
 

t/t1029-mf-priority-default-bank.t

@@ -0,0 +1,99 @@
+#!/bin/bash
+
+test_description='test adding default bank to jobspec in multi-factor priority plugin'
+
+. `dirname $0`/sharness.sh
+MULTI_FACTOR_PRIORITY=${FLUX_BUILD_DIR}/src/plugins/.libs/mf_priority.so
+SUBMIT_AS=${SHARNESS_TEST_SRCDIR}/scripts/submit_as.py
+DB_PATH=$(pwd)/FluxAccountingTest.db
+
+export TEST_UNDER_FLUX_NO_JOB_EXEC=y
+export TEST_UNDER_FLUX_SCHED_SIMPLE_MODE="limited=1"
+test_under_flux 1 job
+
+flux setattr log-stderr-level 1
+
+test_expect_success 'load multi-factor priority plugin' '
+	flux jobtap load -r .priority-default ${MULTI_FACTOR_PRIORITY}
+'
+
+test_expect_success 'check that mf_priority plugin is loaded' '
+	flux jobtap list | grep mf_priority
+'
+
+test_expect_success 'create flux-accounting DB' '
+	flux account -p $(pwd)/FluxAccountingTest.db create-db
+'
+
+test_expect_success 'start flux-accounting service' '
+	flux account-service -p ${DB_PATH} -t
+'
+
+test_expect_success 'add some banks to the DB' '
+	flux account -p ${DB_PATH} add-bank root 1 &&
+	flux account -p ${DB_PATH} add-bank --parent-bank=root account1 1 &&
+	flux account -p ${DB_PATH} add-bank --parent-bank=root account2 1
+'
+
+test_expect_success 'add a user to the DB' '
+	flux account -p ${DB_PATH} add-user --username=user1001 --userid=1001 --bank=account1 &&
+	flux account -p ${DB_PATH} add-user --username=user1001 --userid=1001 --bank=account2
+'
+
+test_expect_success 'submit a job before plugin has any flux-accounting information' '
+	jobid=$(flux python ${SUBMIT_AS} 1002 hostname) &&
+	flux job wait-event -vt 60 $jobid depend &&
+	flux job info $jobid eventlog > eventlog.out &&
+	grep "depend" eventlog.out
+'
+
+test_expect_success 'add user to flux-accounting DB and update plugin' '
+	flux account -p ${DB_PATH} add-user --username=user1002 --userid=1002 --bank=account1 &&
+	flux account-priority-update -p $(pwd)/FluxAccountingTest.db
+'
+
+test_expect_success 'check that bank was added to jobspec' '
+	flux job wait-event -f json $jobid priority &&
+	flux job info $jobid eventlog > eventlog.out &&
+	grep "{\"attributes.system.bank\":\"account1\"}" eventlog.out &&
+	flux job cancel $jobid
+'
+
+test_expect_success 'send flux-accounting DB information to the plugin' '
+	flux account-priority-update -p $(pwd)/FluxAccountingTest.db
+'
+
+test_expect_success 'successfully submit a job under a specified bank' '
+	jobid=$(flux python ${SUBMIT_AS} 1001 --setattr=system.bank=account2 hostname) &&
+	flux job wait-event -f json $jobid priority &&
+	flux job info $jobid jobspec > jobspec.out &&
+	grep "account2" jobspec.out &&
+	flux job cancel $jobid
+'
+
+test_expect_success 'successfully submit a job under a default bank' '
+	jobid=$(flux python ${SUBMIT_AS} 1001 hostname) &&
+	flux job wait-event -f json $jobid priority &&
+	flux job info $jobid eventlog > eventlog.out &&
+	grep "{\"attributes.system.bank\":\"account1\"}" eventlog.out &&
+	flux job cancel $jobid
+'
+
+test_expect_success 'update the default bank for the user' '
+	flux account -p ${DB_PATH} edit-user user1001 --default-bank=account2 &&
+	flux account-priority-update -p ${DB_PATH}
+'
+
+test_expect_success 'successfully submit a job under the new default bank' '
+	jobid=$(flux python ${SUBMIT_AS} 1001 hostname) &&
+	flux job wait-event -f json $jobid priority &&
+	flux job info $jobid eventlog > eventlog.out &&
+	grep "{\"attributes.system.bank\":\"account2\"}" eventlog.out &&
+	flux job cancel $jobid
+'
+
+test_expect_success 'shut down flux-accounting service' '
+	flux python -c "import flux; flux.Flux().rpc(\"accounting.shutdown_service\").get()"
+'
+
+test_done