Python RefHelper.iteritems方法代码示例

# 需要导入模块: from calico.felix.refcount import RefHelper [as 别名]
# 或者: from calico.felix.refcount.RefHelper import iteritems [as 别名]
class TestRefHelper(TestReferenceManager):
    def setUp(self):
        super(TestRefHelper, self).setUp()
        self._rh = RefHelper(self._rm,
                             self._rm,
                             self._rm.ready_callback)

    def test_no_refs(self):
        # With no references, we're ready but haven't been notified
        self.assertFalse(self._rm._ready_called)
        self.assertTrue(self._rh.ready)

        # Discarding non-existent references is allowed
        self._rh.discard_ref("foo")

    def test_acquire_discard_1(self):
        # Acquire a reference to 'foo' - it won't be ready immediately
        self._rh.acquire_ref("foo")
        self.assertFalse(self._rm._ready_called)
        self.assertFalse(self._rh.ready)

        # Spin the actor framework - we become ready
        _, obj = self.call_via_cb(self._rm.get_and_incref, "bar", async=True)
        self.assertTrue(self._rm._ready_called)
        self.assertTrue(self._rh.ready)
        self.assertEqual(next(self._rh.iteritems())[0], "foo")

        # Acquiring an already-acquired reference is idempotent
        self._rh.acquire_ref("foo")
        self.assertTrue(self._rh.ready)

        # Discard the reference
        self._rh.discard_ref("foo")
        _, obj = self.call_via_cb(self._rm.get_and_incref, "baz", async=True)
        self.assertTrue(self._rh.ready)

    def test_sync_acquire_discard(self):
        # Acquire a reference and discard it before it's become ready
        self._rh.acquire_ref("foo")
        self.assertFalse(self._rh.ready)

        self._rh.discard_ref("foo")
        self.assertTrue(self._rh.ready)

        # Spin the actor framework
        _, obj = self.call_via_cb(self._rm.get_and_incref, "bar", async=True)

    def test_acquire_discard_2(self):
        # Acquire two references
        self._rh.acquire_ref("foo")
        _, obj = self.call_via_cb(self._rm.get_and_incref, "bar", async=True)
        self._rh.acquire_ref("baz")
        self.assertFalse(self._rh.ready)
        _, obj = self.call_via_cb(self._rm.get_and_incref, "bar2", async=True)
        acq_ids = list(key for key, value in self._rh.iteritems())
        self.assertItemsEqual(acq_ids, ["foo", "baz"])
        self.assertTrue(self._rh.ready)

        # Discard them all!
        self._rh.discard_all()

# 需要导入模块: from calico.felix.refcount import RefHelper [as 别名]
# 或者: from calico.felix.refcount.RefHelper import iteritems [as 别名]
class ProfileRules(RefCountedActor):
    """
    Actor that owns the per-profile rules chains.
    """
    def __init__(self, profile_id, ip_version, iptables_updater, ipset_mgr):
        super(ProfileRules, self).__init__(qualifier=profile_id)
        assert profile_id is not None

        self.id = profile_id
        self.ip_version = ip_version
        self.ipset_mgr = ipset_mgr
        self._iptables_updater = iptables_updater
        self.notified_ready = False

        self.ipset_refs = RefHelper(self, ipset_mgr, self._maybe_update)

        self._profile = None
        """
        :type dict|None: filled in by first update.  Reset to None on delete.
        """
        self.dead = False

        self.chain_names = {
            "inbound": profile_to_chain_name("inbound", profile_id),
            "outbound": profile_to_chain_name("outbound", profile_id),
        }
        _log.info("Profile %s has chain names %s",
                  profile_id, self.chain_names)

    @actor_message()
    def on_profile_update(self, profile):
        """
        Update the programmed iptables configuration with the new
        profile.
        """
        _log.debug("%s: Profile update: %s", self, profile)
        assert profile is None or profile["id"] == self.id
        assert not self.dead, "Shouldn't receive updates after we're dead."

        old_tags = extract_tags_from_profile(self._profile)
        new_tags = extract_tags_from_profile(profile)

        removed_tags = old_tags - new_tags
        added_tags = new_tags - old_tags
        for tag in removed_tags:
            _log.debug("Queueing ipset for tag %s for decref", tag)
            self.ipset_refs.discard_ref(tag)
        for tag in added_tags:
            _log.debug("Requesting ipset for tag %s", tag)
            self.ipset_refs.acquire_ref(tag)

        self._profile = profile
        self._maybe_update()

    def _maybe_update(self):
        if self.dead:
            _log.info("Not updating: profile is dead.")
        elif not self.ipset_refs.ready:
            _log.info("Can't program rules %s yet, waiting on ipsets",
                      self.id)
        else:
            _log.info("Ready to program rules for %s", self.id)
            self._update_chains()

    @actor_message()
    def on_unreferenced(self):
        """
        Called to tell us that this profile is no longer needed.  Removes
        our iptables configuration.
        """
        try:
            _log.info("%s unreferenced, removing our chains", self)
            self.dead = True
            chains = []
            for direction in ["inbound", "outbound"]:
                chain_name = self.chain_names[direction]
                chains.append(chain_name)
            self._iptables_updater.delete_chains(chains, async=False)
            self.ipset_refs.discard_all()
            self.ipset_refs = None # Break ref cycle.
            self._profile = None
        finally:
            self._notify_cleanup_complete()

    def _update_chains(self):
        """
        Updates the chains in the dataplane.
        """
        _log.info("%s Programming iptables with our chains.", self)
        updates = {}
        for direction in ("inbound", "outbound"):
            _log.debug("Updating %s chain for profile %s", direction,
                       self.id)
            new_profile = self._profile or {}
            _log.debug("Profile %s: %s", self.id, self._profile)
            rules_key = "%s_rules" % direction
            new_rules = new_profile.get(rules_key, [])
            chain_name = self.chain_names[direction]
            tag_to_ip_set_name = {}
            for tag, ipset in self.ipset_refs.iteritems():
#.........这里部分代码省略.........

# 需要导入模块: from calico.felix.refcount import RefHelper [as 别名]
# 或者: from calico.felix.refcount.RefHelper import iteritems [as 别名]

#.........这里部分代码省略.........
                    self._profile = None
                    self._pending_profile = None
                finally:
                    self._cleaned_up = True
                    self._notify_cleanup_complete()
        else:
            if self._pending_profile != self._profile:
                _log.debug("Profile data changed, updating ipset references.")
                # Make sure that all the new tags and selectors are active.
                # We can't discard unneeded ones until we've updated iptables.
                new_tags_and_sels = extract_tags_and_selectors_from_profile(
                    self._pending_profile
                )
                for tag_or_sel in new_tags_and_sels:
                    _log.debug("Requesting ipset for tag %s", tag_or_sel)
                    # Note: acquire_ref() is a no-op if already acquired.
                    self._ipset_refs.acquire_ref(tag_or_sel)

                self._dirty = True
                self._profile = self._pending_profile
                self._required_ipsets = new_tags_and_sels

            if (self._dirty and
                    self._ipset_refs.ready and
                    self._pending_profile is not None):
                _log.info("Ready to program rules for %s", self.id)
                try:
                    self._update_chains()
                except FailedSystemCall as e:
                    _log.error("Failed to program profile chain %s; error: %r",
                               self, e)
                else:
                    # Now we've updated iptables, we can tell the RefHelper
                    # to discard the tags we no longer need.
                    self._ipset_refs.replace_all(self._required_ipsets)
                    self._dirty = False
            elif not self._dirty:
                _log.debug("No changes to program.")
            elif self._pending_profile is None:
                _log.info("Profile is None, removing our chains")
                try:
                    self._delete_chains()
                except FailedSystemCall:
                    _log.exception("Failed to delete chains for profile %s",
                                   self.id)
                else:
                    self._ipset_refs.discard_all()
                    self._dirty = False
            else:
                assert not self._ipset_refs.ready
                _log.info("Can't program rules %s yet, waiting on ipsets",
                          self.id)

    def _delete_chains(self):
        """
        Removes our chains from the dataplane, blocks until complete.
        """
        # Need to block here: have to wait for chains to be deleted
        # before we can decref our ipsets.
        self._iptables_updater.delete_chains(
            self.iptables_generator.profile_chain_names(self.id),
            async=False)

    def _update_chains(self):
        """
        Updates the chains in the dataplane.

        Blocks until the update is complete.

        On entry, self._pending_profile must not be None.

        :raises FailedSystemCall: if the update fails.
        """
        _log.info("%s Programming iptables with our chains.", self)
        assert self._pending_profile is not None, \
            "_update_chains called with no _pending_profile"
        tag_to_ip_set_name = {}
        sel_to_ip_set_name = {}
        for tag_or_sel, ipset in self._ipset_refs.iteritems():
            if isinstance(tag_or_sel, SelectorExpression):
                sel_to_ip_set_name[tag_or_sel] = ipset.ipset_name
            else:
                tag_to_ip_set_name[tag_or_sel] = ipset.ipset_name

        _log.info("Updating chains for profile %s", self.id)
        _log.debug("Profile %s: %s", self.id, self._profile)

        updates, deps = self.iptables_generator.profile_updates(
            self.id,
            self._pending_profile,
            self.ip_version,
            tag_to_ipset=tag_to_ip_set_name,
            selector_to_ipset=sel_to_ip_set_name,
            on_allow="RETURN",
            comment_tag=self.id)

        _log.debug("Queueing programming for rules %s: %s", self.id,
                   updates)

        self._iptables_updater.rewrite_chains(updates, deps, async=False)

# 需要导入模块: from calico.felix.refcount import RefHelper [as 别名]
# 或者: from calico.felix.refcount.RefHelper import iteritems [as 别名]

#.........这里部分代码省略.........
        Called to tell us that this profile is no longer needed.
        """
        # Flag that we're dead and then let finish_msg_batch() do the cleanup.
        self._dead = True

    def _on_ipsets_acquired(self):
        """
        Callback from the RefHelper once it's acquired all the ipsets we
        need.

        This is called from an actor_message on our greenlet.
        """
        # Nothing to do here, if this is being called then we're already in
        # a message batch so _finish_msg_batch() will get called next.
        _log.info("All required ipsets acquired.")

    def _finish_msg_batch(self, batch, results):
        # Due to dependency management in IptablesUpdater, we don't need to
        # worry about programming the dataplane before notifying so do it on
        # this common code path.
        if not self._notified_ready:
            self._notify_ready()
            self._notified_ready = True

        if self._dead:
            # Only want to clean up once.  Note: we can get here a second time
            # if we had a pending ipset incref in-flight when we were asked
            # to clean up.
            if not self._cleaned_up:
                try:
                    _log.info("%s unreferenced, removing our chains", self)
                    chains = set(self.chain_names.values())
                    # Need to block here: have to wait for chains to be deleted
                    # before we can decref our ipsets.
                    self._iptables_updater.delete_chains(chains, async=False)
                    self._ipset_refs.discard_all()
                    self._ipset_refs = None # Break ref cycle.
                    self._profile = None
                    self._pending_profile = None
                finally:
                    self._cleaned_up = True
                    self._notify_cleanup_complete()
        else:
            if self._pending_profile != self._profile:
                _log.debug("Profile data changed, updating ipset references.")
                old_tags = extract_tags_from_profile(self._profile)
                new_tags = extract_tags_from_profile(self._pending_profile)
                removed_tags = old_tags - new_tags
                added_tags = new_tags - old_tags
                for tag in removed_tags:
                    _log.debug("Queueing ipset for tag %s for decref", tag)
                    self._ipset_refs.discard_ref(tag)
                for tag in added_tags:
                    _log.debug("Requesting ipset for tag %s", tag)
                    self._ipset_refs.acquire_ref(tag)
                self._dirty = True
                self._profile = self._pending_profile

            if self._dirty and self._ipset_refs.ready:
                _log.info("Ready to program rules for %s", self.id)
                try:
                    self._update_chains()
                except CalledProcessError as e:
                    _log.error("Failed to program profile chain %s; error: %r",
                               self, e)
                else:
                    self._dirty = False
            elif not self._dirty:
                _log.debug("No changes to program.")
            elif not self._ipset_refs.ready:
                _log.info("Can't program rules %s yet, waiting on ipsets",
                          self.id)

    def _update_chains(self):
        """
        Updates the chains in the dataplane.
        """
        _log.info("%s Programming iptables with our chains.", self)
        updates = {}
        for direction in ("inbound", "outbound"):
            chain_name = self.chain_names[direction]
            _log.info("Updating %s chain %r for profile %s",
                      direction, chain_name, self.id)
            _log.debug("Profile %s: %s", self.id, self._profile)
            new_profile = self._pending_profile or {}
            rules_key = "%s_rules" % direction
            new_rules = new_profile.get(rules_key, [])
            tag_to_ip_set_name = {}
            for tag, ipset in self._ipset_refs.iteritems():
                tag_to_ip_set_name[tag] = ipset.name
            updates[chain_name] = rules_to_chain_rewrite_lines(
                chain_name,
                new_rules,
                self.ip_version,
                tag_to_ip_set_name,
                on_allow="RETURN",
                comment_tag=self.id)
        _log.debug("Queueing programming for rules %s: %s", self.id,
                   updates)
        self._iptables_updater.rewrite_chains(updates, {}, async=False)