Repository: awslabs/llrt
File: modules/llrt_http/src/cached_dns_resolver.rs (Rust)

Code before:
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, name: Name) -> Self::Future {
        let cache = self.cache.clone();
        let permits = self.concurrency;
        let ttl = self.ttl;

        Box::pin(async move {
            let guard = match cache.get_value_or_guard_async(&name).await {
                Ok(guard) => guard,
                Err(placeholder) => {
                    let guard = CacheConcurrencyGuard::new(permits);
                    _ = placeholder.insert(guard.clone());
                    guard
                },
            };
            if let Some(entry) = guard.entry {
                if entry.ttl > Instant::now() {
                    return Ok(entry.addrs);
                }
            };

            let binding = guard
                .semaphore
                .ok_or_else(|| io::Error::new(ErrorKind::InvalidData, "Missing semaphore"))?;
            let lock = binding.acquire().await.unwrap();

            if let Some(item) = cache.get(&name).and_then(|guard| guard.entry) {
                return Ok(item.addrs);
            }

            let addrs = tokio::net::lookup_host((name.as_str(), 0)).await?;
            let addrs = addrs.collect::<Vec<_>>();
            let addrs = SocketAddrs {
                iter: addrs.into_iter(),
            };
            let addrs2 = addrs.clone();
            let entry = CacheEntry {
                ttl: Instant::now() + ttl,
                addrs,
            };
            cache.insert(
                name,
                CacheConcurrencyGuard {
                    semaphore: None,
                    entry: Some(entry),
                },
            );
            drop(lock);
            Ok(addrs2)

Review comment (performance):
This is unsound. Consider the following:
1. Cache is empty, you ask for domain `example.com`
2. Cache add a `CacheConcurrencyGuard` with a semaphore
3. Semaphore is acquired and domain is resolved
4. Replace the name with resolved entry and semaphore none
5. Time passes (more than TTL)
6. You ask `example.com` again, cache returns element
7. entry ttl is expired (line 80), we continue
8. semaphore is None, return an error

Code after:
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, name: Name) -> Self::Future {
        let cache = self.cache.clone();
        let permits = self.concurrency;
        let ttl = self.ttl;

        Box::pin(async move {
            let guard = match cache.get_value_or_guard_async(&name).await {
                Ok(guard) => guard,
                Err(placeholder) => {
                    let guard = CacheConcurrencyGuard::new(permits);
                    _ = placeholder.insert(guard.clone());
                    guard
                },
            };
            if let Some(entry) = guard.entry {
                if entry.ttl > Instant::now() {
                    return Ok(entry.addrs);
                }
            };

            let semaphore = guard.semaphore;
            let semaphore2 = semaphore.clone();
            let lock = semaphore2.acquire().await.unwrap();

            if let Some(item) = cache.get(&name).and_then(|guard| guard.entry) {
                return Ok(item.addrs);
            }

            let addrs = tokio::net::lookup_host((name.as_str(), 0)).await?;
            let addrs = addrs.collect::<Vec<_>>();
            let addrs = SocketAddrs {
                iter: addrs.into_iter(),
            };
            let addrs2 = addrs.clone();
            let entry = CacheEntry {
                ttl: Instant::now() + ttl,
                addrs,
            };
            cache.insert(
                name,
                CacheConcurrencyGuard {
                    semaphore,
                    entry: Some(entry),
                },
            );
            drop(lock);
            Ok(addrs2)
        })