Repository: LMCache/LMCache
File: lmcache/v1/storage_backend/p2p_backend.py (Python)

Code before:
        self,
        peer_init_url: str,
        force_update: bool = False,
    ) -> None:
        if not force_update and peer_init_url in self.peer_id_to_lookup_url_mapping:
            return
        init_side_msg = P2PInitSideMsg()
        init_ret_msg = await self.transfer_channel.async_lazy_init_peer_connection(
            local_id=self.peer_init_url,
            peer_id=peer_init_url,
            peer_init_url=peer_init_url,
            init_side_msg=init_side_msg,
        )
        assert isinstance(init_ret_msg, P2PInitSideRetMsg)
        # update peer_lookup_url
        new_peer_lookup_url = init_ret_msg.peer_lookup_url
        old_peer_lookup_url = self.peer_id_to_lookup_url_mapping.pop(
            peer_init_url, None
        )
        self.peer_id_to_lookup_url_mapping[peer_init_url] = new_peer_lookup_url
        if (
            old_peer_lookup_url is not None
            and old_peer_lookup_url != new_peer_lookup_url
        ):
            logger.info(
                "Peer %s lookup URL changed from %s to %s",
                peer_init_url,
                old_peer_lookup_url,
                new_peer_lookup_url,
            )

        # update lookup_socket, lookup_socket must be updated
        if old_peer_lookup_url is not None:
            old_lookup_socket = self.lookup_url_to_socket_mapping.pop(
                old_peer_lookup_url, None
            )
            if old_lookup_socket is not None:
                try:
                    old_lookup_socket.close(linger=0)
                except Exception as e:
                    logger.error("Error closing old lookup socket", e)
        new_lookup_socket = get_zmq_socket_with_timeout(
            self.async_context,
            new_peer_lookup_url,
            "tcp",
            zmq.REQ,
            "connect",
            self.socket_recv_timeout_ms,
            self.socket_send_timeout_ms,
        )
        self.lookup_url_to_socket_mapping[new_peer_lookup_url] = new_lookup_socket

Review comment (refactor):
The following code discusses the case of `old_peer_lookup_url is not None` in multiple places. To simplify logic, is it possible to perform only one
`if old_peer_lookup_url is not None:`
at the beginning, and recycle / remove all previous locks inside the code block?

Code after:

            return BatchedLookupAndPutRetMsg(num_read_chunks=len(local_mem_objs))
        except Exception as e:
            logger.error(
                "Error during P2P batched lookup and put operation: %s",
                e,
                exc_info=True,
            )
            return BatchedLookupAndPutRetMsg(num_read_chunks=0)

    async def _ensure_peer_connection(
        self,
        peer_init_url: str,
        force_update: bool = False,
    ) -> None:
        if not force_update and peer_init_url in self.peer_id_to_lookup_url_mapping:
            return
        init_side_msg = P2PInitSideMsg()
        init_ret_msg = await self.transfer_channel.async_lazy_init_peer_connection(
            local_id=self.peer_init_url,
            peer_id=peer_init_url,
            peer_init_url=peer_init_url,
            init_side_msg=init_side_msg,
        )
        assert isinstance(init_ret_msg, P2PInitSideRetMsg)
        # update peer_lookup_url
        new_peer_lookup_url = init_ret_msg.peer_lookup_url
        old_peer_lookup_url = self.peer_id_to_lookup_url_mapping.pop(
            peer_init_url, None
        )
        self.peer_id_to_lookup_url_mapping[peer_init_url] = new_peer_lookup_url
        if old_peer_lookup_url is not None:
            if old_peer_lookup_url != new_peer_lookup_url:
                logger.info(
                    "Peer %s lookup URL changed from %s to %s",
                    peer_init_url,
                    old_peer_lookup_url,
                    new_peer_lookup_url,
                )

            # close old lookup socket
            old_lookup_socket = self.lookup_url_to_socket_mapping.pop(
                old_peer_lookup_url, None
            )
            if old_lookup_socket is not None:
                try:
                    old_lookup_socket.close(linger=0)
                except Exception as e:
                    logger.error("Error closing old lookup socket", e)

            # get lookup lock