Skip to content

Commit b4cc92c

Browse files
rohanshah18rohanshah18
authored
Create new config per connection (#178)
## Problem When connecting to multiple indexes using the same Pinecone instance, the connections for each index are cached in a ConcurrentHashMap called connectionsMap, which maps an indexName to its corresponding PineconeConnection. However, after establishing connections for multiple indexes, all entries in the map incorrectly point to the most recently created PineconeConnection instead of maintaining separate connections for each index. ## Solution When assigning a host to the PineconeConfig object, which is passed as input for establishing a gRPC connection, create a new PineconeConfig instance instead of modifying the existing one. Otherwise, all connection objects stored in the connectionsMap will share the same PineconeConfig instance, causing them to reference the most recently created PineconeConnection instead of maintaining separate configurations for each index. ## Type of Change - [X] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] This change requires a documentation update - [ ] Infrastructure change (CI configs, etc) - [ ] Non-code change (docs, etc) - [ ] None of the above: (explain here) ## Test Plan Added integration test.
1 parent 6ed8d46 commit b4cc92c

File tree

3 files changed

+145
-3
lines changed

3 files changed

+145
-3
lines changed

‎src/integration/java/io/pinecone/clients/ConnectionsMapTest.java‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
package io.pinecone.clients;
2+
3+
import io.pinecone.configs.PineconeConnection;
4+
import io.pinecone.exceptions.PineconeNotFoundException;
5+
import io.pinecone.helpers.RandomStringBuilder;
6+
import org.junit.jupiter.api.BeforeAll;
7+
import org.junit.jupiter.api.Test;
8+
import org.openapitools.db_control.client.model.DeletionProtection;
9+
import org.openapitools.db_control.client.model.IndexModel;
10+
11+
import java.util.HashMap;
12+
import java.util.Map;
13+
import java.util.concurrent.ConcurrentHashMap;
14+
15+
import static io.pinecone.helpers.TestUtilities.waitUntilIndexIsReady;
16+
import static org.junit.jupiter.api.Assertions.assertEquals;
17+
import static org.junit.jupiter.api.Assertions.assertThrows;
18+
19+
public class ConnectionsMapTest {
20+
static String indexName1;
21+
static String indexName2;
22+
static Pinecone pinecone1;
23+
static Pinecone pinecone2;
24+
25+
@BeforeAll
26+
public static void setUp() throws InterruptedException {
27+
indexName1 = RandomStringBuilder.build("conn-map-index1", 5);
28+
indexName2 = RandomStringBuilder.build("conn-map-index2", 5);
29+
pinecone1 = new Pinecone
30+
.Builder(System.getenv("PINECONE_API_KEY"))
31+
.withSourceTag("pinecone_test")
32+
.build();
33+
34+
pinecone2 = new Pinecone
35+
.Builder(System.getenv("PINECONE_API_KEY"))
36+
.withSourceTag("pinecone_test")
37+
.build();
38+
}
39+
40+
@Test
41+
public void testMultipleIndexesWithMultipleClients() throws InterruptedException {
42+
Map<String, String> tags = new HashMap<>();
43+
tags.put("env", "test");
44+
45+
// Create index-1
46+
pinecone1.createServerlessIndex(indexName1,
47+
null,
48+
3,
49+
"aws",
50+
"us-east-1",
51+
DeletionProtection.DISABLED,
52+
tags);
53+
54+
// Wait for index to be ready
55+
IndexModel indexModel1 = waitUntilIndexIsReady(pinecone1, indexName1);
56+
57+
// Get index1's host
58+
String host1 = indexModel1.getHost();
59+
60+
// Create index-2
61+
pinecone1.createServerlessIndex(indexName2,
62+
null,
63+
3,
64+
"aws",
65+
"us-east-1",
66+
DeletionProtection.DISABLED,
67+
tags);
68+
69+
// Wait for index to be ready
70+
IndexModel indexModel2 = waitUntilIndexIsReady(pinecone1, indexName2);
71+
72+
// Get index2's host
73+
String host2 = indexModel2.getHost();
74+
75+
// Establish grpc connection for index-1
76+
Index index1_1 = pinecone1.getIndexConnection(indexName1);
77+
// Get connections map
78+
ConcurrentHashMap<String, PineconeConnection> connectionsMap1_1 = pinecone1.getConnectionsMap();
79+
80+
// Verify connectionsMap contains only one <indexName, connection> pair i.e. for index1 and its connection
81+
assertEquals(pinecone1.getConnectionsMap().size(), 1);
82+
// Verify the value for index1 by comparing its value with host1 in the connectionsMap
83+
assertEquals(host1, connectionsMap1_1.get(indexName1).toString());
84+
85+
// Establish grpc connection for index-2
86+
Index index1_2 = pinecone1.getIndexConnection(indexName2);
87+
// Get connections map after establishing second connection
88+
ConcurrentHashMap<String, PineconeConnection> connectionsMap1_2 = pinecone1.getConnectionsMap();
89+
90+
// Verify connectionsMap contains two <indexName, connection> pairs i.e. for index1 and index2
91+
assertEquals(connectionsMap1_2.size(), 2);
92+
// Verify the values by checking host for both indexName1 and indexName2
93+
assertEquals(host1, connectionsMap1_2.get(indexName1).toString());
94+
assertEquals(host2, connectionsMap1_2.get(indexName2).toString());
95+
96+
// Establishing connections with index1 and index2 using another pinecone client
97+
pinecone2.getConnection(indexName1);
98+
ConcurrentHashMap<String, PineconeConnection> connectionsMap2_1 = pinecone1.getConnectionsMap();
99+
// Verify the new connections map is pointing to the same reference
100+
assert connectionsMap2_1 == connectionsMap1_2;
101+
// Verify the size of connections map is still 2 since the connection for index2 was not closed
102+
assertEquals(2, connectionsMap2_1.size());
103+
// Verify the connection value for index1 is host1
104+
assertEquals(host1, connectionsMap2_1.get(indexName1).toString());
105+
106+
pinecone2.getConnection(indexName2);
107+
ConcurrentHashMap<String, PineconeConnection> connectionsMap2_2 = pinecone1.getConnectionsMap();
108+
// Verify the new connections map is pointing to the same reference
109+
assert connectionsMap2_1 == connectionsMap2_2;
110+
// Verify the size of connections map is still 2 since the connections are not closed
111+
assertEquals(2, connectionsMap2_2.size());
112+
// Verify the values by checking host for both indexName1 and indexName2
113+
assertEquals(host1, connectionsMap2_2.get(indexName1).toString());
114+
assertEquals(host2, connectionsMap2_2.get(indexName2).toString());
115+
116+
// Close the connections
117+
index1_1.close();
118+
index1_2.close();
119+
120+
// Verify the map size is now 0
121+
assertEquals(connectionsMap1_1.size(), 0);
122+
assertEquals(connectionsMap1_2.size(), 0);
123+
assertEquals(connectionsMap2_1.size(), 0);
124+
assertEquals(connectionsMap2_2.size(), 0);
125+
126+
// Delete the indexes
127+
pinecone1.deleteIndex(indexName1);
128+
pinecone1.deleteIndex(indexName2);
129+
130+
// Wait for indexes to be deleted
131+
Thread.sleep(5000);
132+
133+
// Confirm the indexes are deleted by calling describe index which should return resource not found
134+
assertThrows(PineconeNotFoundException.class, () -> pinecone1.describeIndex(indexName1));
135+
assertThrows(PineconeNotFoundException.class, () -> pinecone1.describeIndex(indexName2));
136+
}
137+
}

‎src/main/java/io/pinecone/clients/Pinecone.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -949,7 +949,6 @@ public Index getIndexConnection(String indexName) throws PineconeValidationExcep
949949
throw new PineconeValidationException("Index name cannot be null or empty");
950950
}
951951

952-
config.setHost(getIndexHost(indexName));
953952
PineconeConnection connection = getConnection(indexName);
954953
return new Index(connection, indexName);
955954
}
@@ -979,7 +978,6 @@ public AsyncIndex getAsyncIndexConnection(String indexName) throws PineconeValid
979978
throw new PineconeValidationException("Index name cannot be null or empty");
980979
}
981980

982-
config.setHost(getIndexHost(indexName));
983981
PineconeConnection connection = getConnection(indexName);
984982
return new AsyncIndex(config, connection, indexName);
985983
}
@@ -997,7 +995,9 @@ public Inference getInferenceClient() {
997995
}
998996

999997
PineconeConnection getConnection(String indexName) {
1000-
return connectionsMap.computeIfAbsent(indexName, key -> new PineconeConnection(config));
998+
PineconeConfig perConnectionConfig = new PineconeConfig(config.getApiKey(), config.getSourceTag());
999+
perConnectionConfig.setHost(getIndexHost(indexName));
1000+
return connectionsMap.computeIfAbsent(indexName, key -> new PineconeConnection(perConnectionConfig));
10011001
}
10021002

10031003
ConcurrentHashMap<String, PineconeConnection> getConnectionsMap() {

‎src/main/java/io/pinecone/configs/PineconeConnection.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,4 +203,9 @@ public void close() {
203203
logger.warn("Channel shutdown interrupted before termination confirmed");
204204
}
205205
}
206+
207+
@Override
208+
public String toString() {
209+
return config.getHost();
210+
}
206211
}

0 commit comments

Comments
 (0)