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

Code before:
        )
        self.lookup_url_to_socket_mapping[peer_lookup_url] = lookup_socket
        self.lookup_url_to_lock_mapping[peer_lookup_url] = asyncio.Lock()
        # Store socket parameters for future recreation
        self.lookup_url_to_socket_params[peer_lookup_url] = {
            "peer_init_url": peer_init_url,
        }
        logger.info(
            f"Established connection to peer_init_url {peer_init_url}."
            f" The peer_lookup_url: {peer_lookup_url}"
        )

    def _clear_peer_connection_cache(
        self, peer_init_url: str, registration_time: float
    ) -> None:
        peer_lookup_url, old_registration_time = self.peer_id_to_lookup_url_mapping.pop(
            peer_init_url
        )
        logger.info(
            "clearing peer connection cache for peer_init_url %s, "
            "old registration time: %s, new registration time: %s",
            peer_init_url,
            old_registration_time,
            registration_time,
        )
        self.lookup_url_to_lock_mapping.pop(peer_lookup_url, None)
        self.lookup_url_to_socket_params.pop(peer_lookup_url, None)
        old_socket = self.lookup_url_to_socket_mapping.pop(peer_lookup_url, None)
        if old_socket is not None:
            try:
                old_socket.close(linger=0)
            except Exception as e:
                logger.warning("Failed to close old socket: %s", e)

    async def batched_get_non_blocking(
        self,
        lookup_id: str,
        keys: list[CacheEngineKey],
        transfer_spec: Any = None,
    ) -> list[MemoryObj]:
        peer_init_url, peer_lookup_url, location = self.lookup_id_to_peer_mapping.pop(
            lookup_id
        )
        lookup_socket = self.lookup_url_to_socket_mapping[peer_lookup_url]
        lookup_lock = self.lookup_url_to_lock_mapping[peer_lookup_url]

        assert isinstance(transfer_spec, dict)
        cum_chunk_lengths = transfer_spec.get("cum_chunk_lengths", None)
        assert cum_chunk_lengths is not None, "cum_chunk_lengths must be provided"
        assert isinstance(cum_chunk_lengths, list), "cum_chunk_lengths must be a list"


Review comment (suggestion):
It would be better to maintain a PeerStatus or PeerInfo data class in future, something like
```python
class PeerInfo:
  peer_lock: asyncio.Lock,
  peer_socket: zmq.asyncio.Socket,
  peer_socket_params: dict
```

After that, we can pop these related together.

Code after:

            # get lookup lock
            lookup_lock = self.lookup_url_to_lock_mapping.pop(
                old_peer_lookup_url, asyncio.Lock()
            )
        else:
            lookup_lock = asyncio.Lock()

        # update lookup_lock, note that we don't need to create a new lookup_lock
        # even if the lookup_url changes
        self.lookup_url_to_lock_mapping[new_peer_lookup_url] = lookup_lock
        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
        logger.info(
            f"Established connection to peer_init_url {peer_init_url}."
            f" The peer_lookup_url: {new_peer_lookup_url}"
        )

    async def batched_get_non_blocking(
        self,
        lookup_id: str,
        keys: list[CacheEngineKey],
        transfer_spec: Any = None,
    ) -> list[MemoryObj]:
        peer_init_url, _, _ = self.lookup_id_to_peer_mapping.pop(lookup_id)

        assert isinstance(transfer_spec, dict)
        cum_chunk_lengths = transfer_spec.get("cum_chunk_lengths", None)
        assert cum_chunk_lengths is not None, "cum_chunk_lengths must be provided"
        assert isinstance(cum_chunk_lengths, list), "cum_chunk_lengths must be a list"

        mem_objs = []
        str_keys = []
        for idx, key in enumerate(keys):
            shape = self.full_size_shape.copy()
            shape[self.fmt.token_dim()] = (
                cum_chunk_lengths[idx + 1] - cum_chunk_lengths[idx]
            )
            mem_obj = self.local_cpu_backend.allocate(
                torch.Size(shape), self.dtype, self.fmt
            )
            mem_objs.append(mem_obj)
            str_keys.append(key.to_string())