Use DefinitionContainer to maintain concurrentcy invariants

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/api/types/semtype/Definition.java

@@ -18,17 +18,37 @@
 
 package io.ballerina.runtime.api.types.semtype;
 
+import io.ballerina.runtime.internal.types.semtype.DefinitionContainer;
+
 /**
  * Represent a type definition which will act as a layer of indirection between {@code Env} and the type descriptor.
  *
  * @since 2201.11.0
  */
-public interface Definition {
+public abstract class Definition {
+
+    private DefinitionContainer<? extends Definition> container;
 
     /**
      * Get the {@code SemType} of this definition in the given environment.
      *
      * @param env type environment
      */
-    SemType getSemType(Env env);
+    public abstract SemType getSemType(Env env);
+
+    /**
+     * Register the container as the holder of this definition. Used to maintain concurrency invariants.
+     *
+     * @param container holder of the definition
+     * @see io.ballerina.runtime.internal.types.semtype.DefinitionContainer
+     */
+    public void registerContainer(DefinitionContainer<? extends Definition> container) {
+        this.container = container;
+    }
+
+    protected void notifyContainer() {
+        if (container != null) {
+            container.definitionUpdated();
+        }
+    }
 }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BArrayType.java

@@ -232,7 +232,7 @@ public SemType createSemType() {
         if (defn.isDefinitionReady()) {
             return defn.getSemType(env);
         }
-        var result = defn.setDefinition(ListDefinition::new);
+        var result = defn.trySetDefinition(ListDefinition::new);
         if (!result.updated()) {
             return defn.getSemType(env);
         }
@@ -294,7 +294,7 @@ public Optional<SemType> acceptedTypeOf(Context cx) {
         if (acceptedTypeDefn.isDefinitionReady()) {
             return Optional.of(acceptedTypeDefn.getSemType(cx.env));
         }
-        var result = acceptedTypeDefn.setDefinition(ListDefinition::new);
+        var result = acceptedTypeDefn.trySetDefinition(ListDefinition::new);
         if (!result.updated()) {
             return Optional.of(acceptedTypeDefn.getSemType(env));
         }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BFunctionType.java

@@ -232,7 +232,7 @@ public SemType createSemType() {
         if (defn.isDefinitionReady()) {
             return defn.getSemType(env);
         }
-        var result = defn.setDefinition(FunctionDefinition::new);
+        var result = defn.trySetDefinition(FunctionDefinition::new);
         if (!result.updated()) {
             return defn.getSemType(env);
         }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BMapType.java

@@ -191,7 +191,7 @@ public SemType createSemType() {
         if (defn.isDefinitionReady()) {
             return defn.getSemType(env);
         }
-        var result = defn.setDefinition(MappingDefinition::new);
+        var result = defn.trySetDefinition(MappingDefinition::new);
         if (!result.updated()) {
             return defn.getSemType(env);
         }
@@ -237,7 +237,7 @@ public synchronized Optional<SemType> acceptedTypeOf(Context cx) {
         if (acceptedTypeDefn.isDefinitionReady()) {
             return Optional.of(acceptedTypeDefn.getSemType(env));
         }
-        var result = acceptedTypeDefn.setDefinition(MappingDefinition::new);
+        var result = acceptedTypeDefn.trySetDefinition(MappingDefinition::new);
         if (!result.updated()) {
             return Optional.of(acceptedTypeDefn.getSemType(env));
         }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BObjectType.java

@@ -291,7 +291,7 @@ public final SemType createSemType() {
         if (defn.isDefinitionReady()) {
             innerType = defn.getSemType(env);
         } else {
-            var result = defn.setDefinition(ObjectDefinition::new);
+            var result = defn.trySetDefinition(ObjectDefinition::new);
             if (!result.updated()) {
                 innerType = defn.getSemType(env);
             } else {
@@ -393,7 +393,7 @@ public final Optional<SemType> acceptedTypeOf(Context cx) {
         if (acceptedTypeDefn.isDefinitionReady()) {
             innerType = acceptedTypeDefn.getSemType(env);
         } else {
-            var result = acceptedTypeDefn.setDefinition(ObjectDefinition::new);
+            var result = acceptedTypeDefn.trySetDefinition(ObjectDefinition::new);
             if (!result.updated()) {
                 innerType = acceptedTypeDefn.getSemType(env);
             } else {

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BRecordType.java

@@ -250,7 +250,7 @@ public SemType createSemType() {
         if (defn.isDefinitionReady()) {
             return defn.getSemType(env);
         }
-        var result = defn.setDefinition(MappingDefinition::new);
+        var result = defn.trySetDefinition(MappingDefinition::new);
         if (!result.updated()) {
             return defn.getSemType(env);
         }
@@ -393,7 +393,7 @@ public Optional<SemType> acceptedTypeOf(Context cx) {
         if (acceptedTypeDefn.isDefinitionReady()) {
             return Optional.of(acceptedTypeDefn.getSemType(env));
         }
-        var result = acceptedTypeDefn.setDefinition(MappingDefinition::new);
+        var result = acceptedTypeDefn.trySetDefinition(MappingDefinition::new);
         if (!result.updated()) {
             return Optional.of(acceptedTypeDefn.getSemType(env));
         }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BStreamType.java

@@ -153,7 +153,7 @@ public SemType createSemType() {
         if (definition.isDefinitionReady()) {
             return definition.getSemType(env);
         }
-        var result = definition.setDefinition(StreamDefinition::new);
+        var result = definition.trySetDefinition(StreamDefinition::new);
         if (!result.updated()) {
             return definition.getSemType(env);
         }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BTupleType.java

@@ -336,7 +336,7 @@ public SemType createSemType() {
         if (defn.isDefinitionReady()) {
             return defn.getSemType(env);
         }
-        var result = defn.setDefinition(ListDefinition::new);
+        var result = defn.trySetDefinition(ListDefinition::new);
         if (!result.updated()) {
             return defn.getSemType(env);
         }
@@ -405,7 +405,7 @@ public Optional<SemType> acceptedTypeOf(Context cx) {
         if (acceptedTypeDefn.isDefinitionReady()) {
             return Optional.of(acceptedTypeDefn.getSemType(env));
         }
-        var result = acceptedTypeDefn.setDefinition(ListDefinition::new);
+        var result = acceptedTypeDefn.trySetDefinition(ListDefinition::new);
         if (!result.updated()) {
             return Optional.of(acceptedTypeDefn.getSemType(env));
         }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/semtype/DefinitionContainer.java

@@ -5,14 +5,35 @@
 import io.ballerina.runtime.api.types.semtype.SemType;
 
 import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.function.Supplier;
 
+/**
+ * Container used to maintain concurrency invariants when creating a potentially recursive semtype.
+ *
+ * It maintains fallowing invariants.
+ * 1. When the type is being defined only the thread that is defining the type may proceed
+ * 2. After definition is completed any number of threads may proceed concurrently In order to achieve this container
+ * has three phases (init, defining, defined). At init phase (that is no definition has been set) any number of threads
+ * may proceed concurrently. When a thread sets a definition {@code setDefinition} container enters the defining phase.
+ * In that phase only that thread may continue in {@code getSemType} method (this is to allow for recursive type
+ * definitions). Container registers with the {@code Definition} using {@code registerContainer} method. When the
+ * {@code Definition} has been defined (ie. {@code Env} has an atom corresponding to the definition) it must notify the
+ * container using {@code definitionUpdated} method. At this point container moves to defined phase allowing concurrent
+ * access to {@code getSemType}.
+ *
+ * @param <E> type of the definition
+ * @since 2201.11.0
+ */
 public class DefinitionContainer<E extends Definition> {
 
     private final ReadWriteLock rwLock = new ReentrantReadWriteLock();
     private volatile E definition;
 
+    private final ReentrantLock recTypeLock = new ReentrantLock();
+    private volatile boolean isDefining = false;
+
     public boolean isDefinitionReady() {
         try {
             rwLock.readLock().lock();
@@ -25,23 +46,32 @@ public boolean isDefinitionReady() {
     public SemType getSemType(Env env) {
         try {
             rwLock.readLock().lock();
+            // We don't need this check to be synchronized since {@code trySetDefinition} will hold the write lock until
+            // it completes, So isDefining should always be at a consistent state
+            if (isDefining) {
+                // This should prevent threads other than the defining thread to access the rec atom.
+                recTypeLock.lock();
+            }
             return definition.getSemType(env);
         } finally {
             rwLock.readLock().unlock();
         }
     }
 
-    public DefinitionUpdateResult<E> setDefinition(Supplier<E> supplier) {
+    public DefinitionUpdateResult<E> trySetDefinition(Supplier<E> supplier) {
         try {
             rwLock.writeLock().lock();
             boolean updated;
             E newDefinition;
             if (this.definition != null) {
                 updated = false;
-                newDefinition = this.definition;
+                newDefinition = null;
             } else {
                 updated = true;
                 newDefinition = supplier.get();
+                newDefinition.registerContainer(this);
+                this.recTypeLock.lock();
+                isDefining = true;
                 this.definition = newDefinition;
             }
             return new DefinitionUpdateResult<>(newDefinition, updated);
@@ -53,12 +83,28 @@ public DefinitionUpdateResult<E> setDefinition(Supplier<E> supplier) {
     public void clear() {
         try {
             rwLock.writeLock().lock();
+            // This shouldn't happen because defining thread should hold the lock.
+            assert !isDefining;
             this.definition = null;
         } finally {
             rwLock.writeLock().unlock();
         }
     }
 
+    public void definitionUpdated() {
+        recTypeLock.unlock();
+        isDefining = false;
+    }
+
+    /**
+     * Result of trying to update the definition.
+     *
+     * @param <E>        Type of the definition
+     * @param updated    If update was successful. If this failed you must get the semtype using the {@code getSemType}
+     *                   method of the container
+     * @param definition If update was successful this will be the new definition. Otherwise, this will be null
+     * @since 2201.11.0
+     */
     public record DefinitionUpdateResult<E>(E definition, boolean updated) {
 
     }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/semtype/FunctionDefinition.java

@@ -32,10 +32,10 @@
  *
  * @since 2201.11.0
  */
-public class FunctionDefinition implements Definition {
+public class FunctionDefinition extends Definition {
 
-    private RecAtom rec;
-    private SemType semType;
+    private volatile RecAtom rec;
+    private volatile SemType semType;
 
     @Override
     public SemType getSemType(Env env) {
@@ -65,7 +65,9 @@ public SemType define(Env env, SemType args, SemType ret, FunctionQualifiers qua
         } else {
             atom = env.functionAtom(atomicType);
         }
-        return this.createSemType(atom);
+        SemType semType = this.createSemType(atom);
+        notifyContainer();
+        return semType;
     }
 
 }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/semtype/ListDefinition.java

@@ -39,10 +39,10 @@
  *
  * @since 2201.11.0
  */
-public class ListDefinition implements Definition {
+public class ListDefinition extends Definition {
 
-    private RecAtom rec = null;
-    private SemType semType = null;
+    private volatile RecAtom rec = null;
+    private volatile SemType semType = null;
 
     @Override
     public SemType getSemType(Env env) {
@@ -63,7 +63,9 @@ public SemType defineListTypeWrapped(Env env, SemType[] initial, int fixedLength
         }
         SemType restCell =
                 Builder.getCellContaining(env, union(rest, getUndefType()), isNever(rest) ? CELL_MUT_NONE : mut);
-        return define(env, initialCells, fixedLength, restCell);
+        SemType semType = define(env, initialCells, fixedLength, restCell);
+        notifyContainer();
+        return semType;
     }
 
     private SemType define(Env env, SemType[] initial, int fixedLength, SemType rest) {

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/semtype/MappingDefinition.java

@@ -42,10 +42,10 @@
  *
  * @since 2201.11.0
  */
-public class MappingDefinition implements Definition {
+public class MappingDefinition extends Definition {
 
-    private RecAtom rec = null;
-    private SemType semType = null;
+    private volatile RecAtom rec = null;
+    private volatile SemType semType = null;
 
     @Override
     public SemType getSemType(Env env) {
@@ -75,7 +75,9 @@ public SemType defineMappingTypeWrapped(Env env, Field[] fields, SemType rest, C
         }
         SemType restCell = Builder.getCellContaining(env, union(rest, getUndefType()),
                 isNever(rest) ? CellAtomicType.CellMutability.CELL_MUT_NONE : mut);
-        return define(env, cellFields, restCell);
+        SemType semType = define(env, cellFields, restCell);
+        notifyContainer();
+        return semType;
     }
 
     SemType define(Env env, BCellField[] cellFields, SemType rest) {

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/semtype/ObjectDefinition.java

@@ -46,7 +46,7 @@
  *
  * @since 2201.11.0
  */
-public class ObjectDefinition implements Definition {
+public class ObjectDefinition extends Definition {
 
     private final MappingDefinition mappingDefinition = new MappingDefinition();
 
@@ -63,7 +63,9 @@ public SemType define(Env env, ObjectQualifiers qualifiers, List<Member> members
         SemType mappingType = mappingDefinition.define(env, Stream.concat(memberStream, qualifierStream)
                 .map(field -> MappingDefinition.BCellField.from(env, field, mut))
                 .toArray(MappingDefinition.BCellField[]::new), restMemberType(env, mut, qualifiers.readonly()));
-        return objectContaining(mappingType);
+        SemType semType = objectContaining(mappingType);
+        notifyContainer();
+        return semType;
     }
 
     private SemType objectContaining(SemType mappingType) {

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/semtype/StreamDefinition.java

@@ -35,7 +35,7 @@
  *
  * @since 2201.11.0
  */
-public class StreamDefinition implements Definition {
+public class StreamDefinition extends Definition {
 
     private final ListDefinition listDefinition = new ListDefinition();
 
@@ -50,7 +50,9 @@ public SemType define(Env env, SemType valueType, SemType completionType) {
         }
         SemType tuple = listDefinition.defineListTypeWrapped(env, new SemType[]{valueType, completionType}, 2,
                 Builder.getNeverType(), CellAtomicType.CellMutability.CELL_MUT_LIMITED);
-        return streamContaining(tuple);
+        SemType semType = streamContaining(tuple);
+        notifyContainer();
+        return semType;
     }
 
     private SemType streamContaining(SemType tupleType) {