From 297f6169feecd20e121d102e1b63a505c8b3e74a Mon Sep 17 00:00:00 2001
From: Sage Weil <sage@inktank.com>
Date: Wed, 28 May 2014 13:49:52 -0700
Subject: [PATCH] crush: add get_rule_weight_map

Calculate a weight map of OSDs for a given rule.

Signed-off-by: Sage Weil <sage@inktank.com>
---
 src/crush/CrushWrapper.cc          | 46 ++++++++++++++++++++++++++++++
 src/crush/CrushWrapper.h           | 12 ++++++++
 src/test/crush/TestCrushWrapper.cc |  5 ++++
 3 files changed, 63 insertions(+)

diff --git a/src/crush/CrushWrapper.cc b/src/crush/CrushWrapper.cc
index 52730a1d3903f..43a525a1b1b9c 100644
--- a/src/crush/CrushWrapper.cc
+++ b/src/crush/CrushWrapper.cc
@@ -794,6 +794,52 @@ int CrushWrapper::add_simple_ruleset(string name, string root_name,
   return rno;
 }
 
+int CrushWrapper::get_rule_weight_map(unsigned ruleno, map<int,float> *pmap)
+{
+  if (ruleno >= crush->max_rules)
+    return -ENOENT;
+  if (crush->rules[ruleno] == NULL)
+    return -ENOENT;
+  crush_rule *rule = crush->rules[ruleno];
+
+  // build a weight map for each TAKE in the rule, and then merge them
+  for (unsigned i=0; i<rule->len; ++i) {
+    map<int,float> m;
+    float sum = 0;
+    if (rule->steps[i].op == CRUSH_RULE_TAKE) {
+      int n = rule->steps[i].arg1;
+      if (n >= 0) {
+	m[n] = 1.0;
+	sum = 1.0;
+      } else {
+	list<int> q;
+	q.push_back(n);
+	while (!q.empty()) {
+	  int bno = q.front();
+	  q.pop_front();
+	  crush_bucket *b = crush->buckets[-1-bno];
+	  assert(b);
+	  for (unsigned j=0; j<b->size; ++j) {
+	    float w = crush_get_bucket_item_weight(b, j);
+	    m[b->items[j]] = w;
+	    sum += w;
+	  }
+	}
+      }
+    }
+    for (map<int,float>::iterator p = m.begin(); p != m.end(); ++p) {
+      map<int,float>::iterator q = pmap->find(p->first);
+      if (q == pmap->end()) {
+	(*pmap)[p->first] = p->second / sum;
+      } else {
+	q->second += p->second / sum;
+      }
+    }
+  }
+
+  return 0;
+}
+
 int CrushWrapper::remove_rule(int ruleno)
 {
   if (ruleno >= (int)crush->max_rules)
diff --git a/src/crush/CrushWrapper.h b/src/crush/CrushWrapper.h
index 0a633a51f5cba..bc7bf9b9c91d4 100644
--- a/src/crush/CrushWrapper.h
+++ b/src/crush/CrushWrapper.h
@@ -631,6 +631,18 @@ public:
     return s->arg2;
   }
 
+  /**
+   * calculate a map of osds to weights for a given rule
+   *
+   * Generate a map of which OSDs get how much relative weight for a
+   * given rule.
+   *
+   * @param ruleno [in] rule id
+   * @param pmap [out] map of osd to weight
+   * @return 0 for success, or negative error code
+   */
+  int get_rule_weight_map(unsigned ruleno, map<int,float> *pmap);
+
   /* modifiers */
   int add_rule(int len, int ruleset, int type, int minsize, int maxsize, int ruleno) {
     if (!crush) return -ENOENT;
diff --git a/src/test/crush/TestCrushWrapper.cc b/src/test/crush/TestCrushWrapper.cc
index d70a52540f598..7d42da5030241 100644
--- a/src/test/crush/TestCrushWrapper.cc
+++ b/src/test/crush/TestCrushWrapper.cc
@@ -538,6 +538,11 @@ TEST(CrushWrapper, dump_rules) {
 	      ss.str().find("<item_name>default</item_name></step>"));
   }
 
+  map<int,float> wm;
+  c->get_rule_weight_map(0, &wm);
+  ASSERT_TRUE(wm.size() == 1);
+  ASSERT_TRUE(wm[0] == 1.0);
+
   delete c;
 }
 
-- 
2.39.5