`GoogleGenAiTextEmbeddingModel.dimensions()` mutates a shared static `HashMap`, throwing `ConcurrentModificationException` under concurrent use

[hvmedx]

Hello,

GoogleGenAiTextEmbeddingModel.dimensions() mutates a private static final plain HashMap via computeIfAbsent, and the mapping function performs a network round trip. Concurrent first-time callers for a model that is not in the pre-seeded enum therefore throw ConcurrentModificationException.

I found this because vector searches were intermittently failing on a freshly started service, roughly once per run, always in the first few seconds and never afterwards:

java.util.ConcurrentModificationException: null
	at java.base/java.util.HashMap.computeIfAbsent(HashMap.java:1229)
	at org.springframework.ai.google.genai.text.GoogleGenAiTextEmbeddingModel.dimensions(GoogleGenAiTextEmbeddingModel.java:262)
	at org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore.createObservationContextBuilder(ElasticsearchVectorStore.java:350)
	at org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore.similaritySearch(AbstractObservationVectorStore.java:132)
	at ...application code...

Every frame in that trace lines up exactly with the source, which is worth stating up front since races are often hard to pin down: HashMap.java:1229 in JDK 21 is precisely if (mc != modCount) { throw new ConcurrentModificationException(); }, immediately after V v = mappingFunction.apply(key); on line 1228. The other three frames are the computeIfAbsent call, the .dimensions(...) observation tag, and the createObservationContextBuilder call inside similaritySearch.

The code on main today:

private static final Map<String, Integer> KNOWN_EMBEDDING_DIMENSIONS = Stream
    .of(GoogleGenAiTextEmbeddingModelName.values())
    .collect(Collectors.toMap(GoogleGenAiTextEmbeddingModelName::getName,
            GoogleGenAiTextEmbeddingModelName::getDimensions));

@Override
public int dimensions() {
    return KNOWN_EMBEDDING_DIMENSIONS.computeIfAbsent(this.options.getModel(), model -> super.dimensions());
}

Three things combine to make this reachable rather than theoretical.

The map is a shared static HashMap. The two-argument Collectors.toMap returns a plain HashMap, and the field is static, so it is shared across every instance in the JVM. computeIfAbsent writes to it.

The mapping function always performs a network call. super.dimensions() reaches AbstractEmbeddingModel.dimensions(this, "Test", "Hello World"). Because the model name passed is the hardcoded literal "Test", the lookup in embedding-model-dimensions.properties can never match, so it always falls through to embeddingModel.embed("Hello World"). That holds the map in a mid-mutation state for the length of a full API round trip, hundreds of milliseconds, rather than nanoseconds. (This is the same hardcoded "Test" noted in #2426.)

Any model outside the enum misses on every call until one thread finally inserts. gemini-embedding-001 was added to GoogleGenAiTextEmbeddingModelName in 1.1.4, which incidentally hides the race for the current default, but it remains live for anything the enum has not caught up with: preview and experimental names such as gemini-embedding-exp-03-07, newer text-embedding-00x releases, regional or versioned names, custom endpoints. I originally hit it on 1.1.1, where gemini-embedding-001 itself was not yet listed.

Two further notes on scope and severity.

AbstractObservationVectorStore builds an observation context on add (L79), delete (L104, L116) and query (L132), so any vector store that puts embeddingModel.dimensions() into that context is exposed on all three operations, not only search. And the exception surfaces from observability tag computation on a path the caller believes is a read, which makes it disproportionately hard to diagnose. In our case Spring AI's tool executor swallowed it into a result string, so what we actually saw was RuntimeException: Failed to parse JSON with no application frames in the stack. That cost a day.

Finally, ConcurrentModificationException is the detected subset of the problem, not the whole of it. This is unsynchronised mutation of shared state: modCount is a non-volatile int read without synchronisation, so the guard is itself racy. In a standalone reproduction, an unpredictable subset of threads threw (2 to 6 of 8 across runs), and the threads that did not throw each executed the non-atomic ++size, leaving size() permanently inconsistent with the actual entry count in 16 of 25 runs. Benign in this particular map, since it is only ever read by key, but it is real corruption rather than merely a thrown exception.

Environment:
Spring AI 2.0.0 (code path verified unchanged on main @ b1384c8)
Spring Boot 4.0.x
Java 21 (Temurin 21.0.4)
Vector store: Elasticsearch, dense_vector with kNN
Embedding model: any name not present in GoogleGenAiTextEmbeddingModelName
Originally observed on 1.1.1 with gemini-embedding-001, before that name was added to the enum
Steps to reproduce

Configure GoogleGenAiTextEmbeddingModel with a model name not listed in GoogleGenAiTextEmbeddingModelName, for example gemini-embedding-exp-03-07
Wire it into any VectorStore that reports dimensions in its observation context, such as ElasticsearchVectorStore
From a cold JVM, issue several similaritySearch calls concurrently. Three was enough for us
Intermittently, one or more throw ConcurrentModificationException from HashMap.computeIfAbsent

It only happens on the first calls. Once the key is present, computeIfAbsent returns on the fast path at HashMap.java:1224 without ever reading modCount, and the problem disappears for the life of the JVM. That self-healing is why it reads as flakiness.

Expected behavior:
dimensions() should be safe to call concurrently. Concurrent first callers should either share one computed value or compute independently, but never throw and never mutate shared state unsynchronised.

Minimal Complete Reproducible example:
This needs no API key. It stubs the network call while keeping the timing window that makes the race visible.

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;

import java.util.List;
import java.util.UUID;
import java.util.concurrent.*;

import org.junit.jupiter.api.Test;
import org.springframework.ai.embedding.*;
import org.springframework.ai.google.genai.embedding.GoogleGenAiEmbeddingConnectionDetails;
import org.springframework.ai.google.genai.text.*;

class GoogleGenAiTextEmbeddingModelConcurrencyTest {

    @Test
    void dimensionsIsNotThreadSafeForModelsOutsideTheKnownList() throws Exception {
        var connection = mock(GoogleGenAiEmbeddingConnectionDetails.class);
        // Unique per run: KNOWN_EMBEDDING_DIMENSIONS is static and never reset, so a fixed name
        // would be cached by the first run and the test would pass on every later run in the same JVM.
        var options = GoogleGenAiTextEmbeddingOptions.builder()
                .model("unknown-model-" + UUID.randomUUID())
                .build();

        var model = spy(new GoogleGenAiTextEmbeddingModel(connection, options));
        // Stands in for the embed() round trip super.dimensions() performs.
        doAnswer(invocation -> {
            Thread.sleep(200);
            return new EmbeddingResponse(List.of(new Embedding(new float[3072], 0)));
        }).when(model).call(any(EmbeddingRequest.class));

        int threads = 8;
        var start = new CountDownLatch(1);
        var failures = new CopyOnWriteArrayList<Throwable>();
        var pool = Executors.newFixedThreadPool(threads);

        for (int i = 0; i < threads; i++) {
            pool.submit(() -> {
                try {
                    start.await();
                    model.dimensions();
                }
                catch (Throwable t) {
                    failures.add(t);
                }
            });
        }
        start.countDown();
        pool.shutdown();
        pool.awaitTermination(30, TimeUnit.SECONDS);

        assertThat(failures).isEmpty();   // fails: java.util.ConcurrentModificationException
    }
}

Being upfront: it is a race, so it is probabilistic. It reproduces consistently for me but you may want a couple of runs. If you would prefer something deterministic, asserting that dimensions() does not mutate shared static state pins the invariant directly rather than the symptom.

Possible fix:
I think the smallest correct change is to stop mutating the shared table and let the base class do the caching it already does:

@Override
public int dimensions() {
    Integer known = KNOWN_EMBEDDING_DIMENSIONS.get(this.options.getModel());
    return (known != null) ? known : super.dimensions();
}

This keeps the benefit of #4803, which introduced the current form specifically to avoid eagerly computing dimensions for known models: known models still short-circuit without touching the network. Unknown models fall through to AbstractEmbeddingModel, which memoises per instance in an AtomicInteger. That memoisation is itself only a non-atomic check-then-set, so two threads can duplicate the API call, but it never throws and never corrupts anything.

Switching the field to ConcurrentHashMap would also stop the exception, but it serialises all concurrent callers behind one network call, and ConcurrentHashMap.computeIfAbsent has its own restrictions on what the mapping function may do. The read-only version above seems cleaner. Making the field Map.copyOf(...) as well would turn any future computeIfAbsent into a loud failure during development rather than a racy one in production.

Two pieces of context that may help you place this:

It is a regression, and it is the only one of its kind in the repo. #4803 (commit a8b3982, first shipped in 1.1.0) changed this method from a read-only lookup to computeIfAbsent. Every sibling embedding model still reads its table without mutating: VertexAiTextEmbeddingModel, which this class was forked from, uses getOrDefault(model, super.dimensions()), and MistralAiEmbeddingModel uses an immutable Map.of(...). This is the only embedding model in the codebase that writes to its known-dimensions table.
An open PR would spread the pattern. #4810 proposes the same computeIfAbsent approach for MistralAiEmbeddingModel. Moving the map to an instance field does not fix it, since embedding models are singleton beans and the concurrency is between requests rather than between instances. Worth resolving the pattern before that lands.

Happy to open a PR with the fix and the test if that is useful. Thanks for your time and for the library. :)