[Python] Cache __hash__, fast_name, and _TypedMetricName on MetricName

sdks/python/apache_beam/metrics/metric_test.py

@@ -57,6 +57,29 @@ def test_basic_metric_name(self):
     self.assertEqual(key.metric.name, 'name1')
     self.assertEqual(key, MetricKey('step1', MetricName('namespace1', 'name1')))
 
+  def test_metric_name_pickle_and_legacy_dict_state(self):
+    import pickle
+    name = MetricName('namespace1', 'name1', urn='urn1', labels={'k': 'v'})
+    _ = hash(name)
+    _ = name.fast_name()
+    unpickled = pickle.loads(pickle.dumps(name))
+    self.assertEqual(unpickled, name)
+    self.assertEqual(hash(unpickled), hash(name))
+    self.assertEqual(unpickled.fast_name(), name.fast_name())
+
+    legacy = MetricName.__new__(MetricName)
+    legacy.__setstate__({
+        'namespace': 'namespace1',
+        'name': 'name1',
+        'urn': 'urn1',
+        'labels': {
+            'k': 'v'
+        },
+    })
+    self.assertEqual(legacy, name)
+    self.assertEqual(hash(legacy), hash(name))
+    self.assertEqual(legacy.fast_name(), name.fast_name())
+
 
 class MetricResultsTest(unittest.TestCase):
   def test_metric_filter_namespace_matching(self):

sdks/python/apache_beam/metrics/metricbase.py

@@ -82,9 +82,11 @@ def __init__(
     self.name = name
     self.urn = urn
     self.labels = labels if labels else {}
+    self._cached_fast_name: Optional[str] = None
+    self._cached_hash: Optional[int] = None
 
   def __eq__(self, other):
-    return (
+    return self is other or (
         self.namespace == other.namespace and self.name == other.name and
         self.urn == other.urn and self.labels == other.labels)
 
@@ -97,17 +99,38 @@ def __str__(self):
           self.namespace, self.name)
 
   def __hash__(self):
-    return hash((self.namespace, self.name, self.urn) +
-                tuple(self.labels.items()))
+    if self._cached_hash is None:
+      self._cached_hash = hash(
+          (self.namespace, self.name, self.urn) +
+          tuple(self.labels.items())) if self.labels else hash(
+              (self.namespace, self.name, self.urn))
+    return self._cached_hash
 
   def fast_name(self):
-    name = self.name or ''
-    namespace = self.namespace or ''
-    urn = self.urn or ''
-    labels = ''
-    if self.labels:
-      labels = '_'.join(['%s=%s' % (k, v) for (k, v) in self.labels.items()])
-    return '%d_%s%s%s%s' % (len(name), name, namespace, urn, labels)
+    if self._cached_fast_name is None:
+      name = self.name or ''
+      namespace = self.namespace or ''
+      urn = self.urn or ''
+      labels = ''
+      if self.labels:
+        labels = '_'.join(['%s=%s' % (k, v) for (k, v) in self.labels.items()])
+      self._cached_fast_name = '%d_%s%s%s%s' % (
+          len(name), name, namespace, urn, labels)
+    return self._cached_fast_name
+
+  def __getstate__(self):
+    return (self.namespace, self.name, self.urn, self.labels)
+
+  def __setstate__(self, state):
+    if isinstance(state, dict):
+      self.namespace = state['namespace']
+      self.name = state['name']
+      self.urn = state['urn']
+      self.labels = state['labels']
+    else:
+      self.namespace, self.name, self.urn, self.labels = state
+    self._cached_fast_name = None
+    self._cached_hash = None
 
 
 class Metric(object):