1 /*
2  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License").
5  * You may not use this file except in compliance with the License.
6  * A copy of the License is located at
7  *
8  *  http://aws.amazon.com/apache2.0
9  *
10  * or in the "license" file accompanying this file. This file is distributed
11  * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12  * express or implied. See the License for the specific language governing
13  * permissions and limitations under the License.
14  */
15 
16 package software.amazon.awssdk.http.nio.netty.internal;
17 
18 import static software.amazon.awssdk.utils.Validate.paramNotNull;
19 
20 import io.netty.channel.pool.ChannelPool;
21 import io.netty.channel.pool.ChannelPoolMap;
22 import java.io.Closeable;
23 import java.util.Collections;
24 import java.util.HashMap;
25 import java.util.Iterator;
26 import java.util.Map;
27 import java.util.concurrent.ConcurrentHashMap;
28 import java.util.concurrent.ConcurrentMap;
29 import software.amazon.awssdk.annotations.SdkInternalApi;
30 
31 /**
32  * Replacement for {@link io.netty.channel.pool.AbstractChannelPoolMap}. This implementation guarantees
33  * only one instance of a {@link ChannelPool} is created for each key.
34  */
35 // TODO do we need to use this for H2?
36 @SdkInternalApi
37 public abstract class SdkChannelPoolMap<K, P extends ChannelPool>
38         implements ChannelPoolMap<K, P>, Iterable<Map.Entry<K, P>>, Closeable {
39 
40     private final ConcurrentMap<K, P> map = new ConcurrentHashMap<>();
41 
42     @Override
get(K key)43     public final P get(K key) {
44         return map.computeIfAbsent(key, this::newPool);
45     }
46 
47     /**
48      * Remove the {@link ChannelPool} from this {@link io.netty.channel.pool.AbstractChannelPoolMap}. Returns {@code true} if
49      * removed, {@code false} otherwise.
50      *
51      * Please note that {@code null} keys are not allowed.
52      */
remove(K key)53     public final boolean remove(K key) {
54         P pool = map.remove(paramNotNull(key, "key"));
55         if (pool != null) {
56             pool.close();
57             return true;
58         }
59         return false;
60     }
61 
62     @Override
iterator()63     public final Iterator<Map.Entry<K, P>> iterator() {
64         return new ReadOnlyIterator<>(map.entrySet().iterator());
65     }
66 
67     /**
68      * Returns the number of {@link ChannelPool}s currently in this {@link io.netty.channel.pool.AbstractChannelPoolMap}.
69      */
size()70     public final int size() {
71         return map.size();
72     }
73 
74     /**
75      * Returns {@code true} if the {@link io.netty.channel.pool.AbstractChannelPoolMap} is empty, otherwise {@code false}.
76      */
isEmpty()77     public final boolean isEmpty() {
78         return map.isEmpty();
79     }
80 
81     @Override
contains(K key)82     public final boolean contains(K key) {
83         return map.containsKey(paramNotNull(key, "key"));
84     }
85 
86     /**
87      * Called once a new {@link ChannelPool} needs to be created as non exists yet for the {@code key}.
88      */
newPool(K key)89     protected abstract P newPool(K key);
90 
91     @Override
close()92     public void close() {
93         map.keySet().forEach(this::remove);
94     }
95 
pools()96     public final Map<K, P> pools() {
97         return Collections.unmodifiableMap(new HashMap<>(map));
98     }
99 
100     /**
101      * {@link Iterator} that prevents removal.
102      *
103      * @param <T> Type of object being iterated.
104      */
105     private final class ReadOnlyIterator<T> implements Iterator<T> {
106         private final Iterator<? extends T> iterator;
107 
ReadOnlyIterator(Iterator<? extends T> iterator)108         private ReadOnlyIterator(Iterator<? extends T> iterator) {
109             this.iterator = paramNotNull(iterator, "iterator");
110         }
111 
112         @Override
hasNext()113         public boolean hasNext() {
114             return this.iterator.hasNext();
115         }
116 
117         @Override
next()118         public T next() {
119             return this.iterator.next();
120         }
121 
122         @Override
remove()123         public void remove() {
124             throw new UnsupportedOperationException("Read-only iterator doesn't support removal.");
125         }
126     }
127 }
128 
129