feat: Restrict cron execution for single pod in clustered environment

app/server/appsmith-server/src/main/java/com/appsmith/server/aspect/FeatureFlaggedMethodInvokerAspect.java

@@ -84,6 +84,7 @@ private Object invokeMethod(Boolean isFeatureSupported, ProceedingJoinPoint join
             if (e instanceof AppsmithException) {
                 throw (AppsmithException) e;
             }
+            log.error("Exception while invoking super class method", e);
             String errorMessage = "Exception while invoking super class method";
             AppsmithException exception = getInvalidAnnotationUsageException(method, errorMessage);
             log.error(exception.getMessage(), e);

app/server/appsmith-server/src/main/java/com/appsmith/server/configurations/InstanceConfig.java

@@ -47,7 +47,7 @@ public void onApplicationEvent(ApplicationReadyEvent applicationReadyEvent) {
         Mono<Void> registrationAndRtsCheckMono = configService
                 .getByName(Appsmith.APPSMITH_REGISTERED)
                 .filter(config -> TRUE.equals(config.getConfig().get("value")))
-                .switchIfEmpty(Mono.defer(() -> instanceConfigHelper.registerInstance()))
+                .switchIfEmpty(Mono.defer(instanceConfigHelper::registerInstance))
                 .onErrorResume(errorSignal -> {
                     log.debug("Instance registration failed with error: \n{}", errorSignal.getMessage());
                     return Mono.empty();
@@ -60,20 +60,12 @@ public void onApplicationEvent(ApplicationReadyEvent applicationReadyEvent) {
                 .flatMap(signal -> registrationAndRtsCheckMono)
                 // Prefill the server cache with anonymous user permission group ids.
                 .then(cacheableRepositoryHelper.preFillAnonymousUserPermissionGroupIdsCache())
-                // Add cold publisher as we have dependency on the instance registration
-                // TODO Update implementation to fetch license status for all the organizations once multi-tenancy is
-                //  introduced
+                // Cold publisher to wait for upstream execution to complete as we have dependency on the instance
+                // registration
                 .then(Mono.defer(instanceConfigHelper::isLicenseValid)
-                        // Ensure that the organization feature flags are refreshed with the latest values after
-                        // completing
-                        // the
+                        // Ensure that the org feature flags are refreshed with the latest values after completing the
                         // license verification process.
-                        .flatMap(isValid -> {
-                            log.debug(
-                                    "License verification completed with status: {}",
-                                    TRUE.equals(isValid) ? "valid" : "invalid");
-                            return instanceConfigHelper.updateCacheForOrganizationFeatureFlags();
-                        }));
+                        .flatMap(isValid -> instanceConfigHelper.updateCacheForOrganizationFeatureFlags()));
 
         try {
             startupProcess.block();

app/server/appsmith-server/src/main/java/com/appsmith/server/domains/Organization.java

@@ -12,6 +12,8 @@
 
 import java.io.Serializable;
 
+import static com.appsmith.external.helpers.StringUtils.dotted;
+
 @Getter
 @Setter
 @ToString
@@ -35,4 +37,9 @@ public class Organization extends BaseDomain implements Serializable {
     OrganizationConfiguration organizationConfiguration;
 
     // TODO add SSO and other configurations here after migrating from environment variables to database configuration
+
+    public static class Fields {
+        public static final String organizationConfiguration_isRestartRequired =
+                dotted(organizationConfiguration, OrganizationConfiguration.Fields.isRestartRequired);
+    }
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/domains/OrganizationConfiguration.java

@@ -3,7 +3,11 @@
 import com.appsmith.server.domains.ce.OrganizationConfigurationCE;
 import lombok.Data;
 import lombok.EqualsAndHashCode;
+import lombok.experimental.FieldNameConstants;
 
 @Data
 @EqualsAndHashCode(callSuper = true)
-public class OrganizationConfiguration extends OrganizationConfigurationCE {}
+@FieldNameConstants
+public class OrganizationConfiguration extends OrganizationConfigurationCE {
+    public static class Fields extends OrganizationConfigurationCE.Fields {}
+}

app/server/appsmith-server/src/main/java/com/appsmith/server/domains/ce/OrganizationConfigurationCE.java

@@ -8,6 +8,7 @@
 import com.appsmith.server.domains.OrganizationConfiguration;
 import com.fasterxml.jackson.annotation.JsonInclude;
 import lombok.Data;
+import lombok.experimental.FieldNameConstants;
 import org.apache.commons.lang3.ObjectUtils;
 
 import java.io.Serializable;
@@ -16,6 +17,7 @@
 import java.util.Map;
 
 @Data
+@FieldNameConstants
 public class OrganizationConfigurationCE implements Serializable {
 
     private String googleMapsKey;
@@ -88,4 +90,6 @@ public void copyNonSensitiveValues(OrganizationConfiguration organizationConfigu
     public Boolean isEmailVerificationEnabled() {
         return Boolean.TRUE.equals(this.emailVerificationEnabled);
     }
+
+    public static class Fields {}
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/helpers/InstanceConfigHelperImpl.java

@@ -7,6 +7,7 @@
 import com.appsmith.server.services.AnalyticsService;
 import com.appsmith.server.services.ConfigService;
 import com.appsmith.server.services.FeatureFlagService;
+import com.appsmith.server.services.OrganizationService;
 import com.appsmith.server.solutions.ReleaseNotesService;
 import org.springframework.context.ApplicationContext;
 import org.springframework.data.mongodb.core.ReactiveMongoTemplate;
@@ -24,7 +25,8 @@ public InstanceConfigHelperImpl(
             AnalyticsService analyticsService,
             NetworkUtils networkUtils,
             ReleaseNotesService releaseNotesService,
-            RTSCaller rtsCaller) {
+            RTSCaller rtsCaller,
+            OrganizationService organizationService) {
         super(
                 configService,
                 cloudServicesConfig,
@@ -35,6 +37,7 @@ public InstanceConfigHelperImpl(
                 analyticsService,
                 networkUtils,
                 releaseNotesService,
-                rtsCaller);
+                rtsCaller,
+                organizationService);
     }
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/helpers/ce/InstanceConfigHelperCEImpl.java

@@ -15,6 +15,7 @@
 import com.appsmith.server.services.AnalyticsService;
 import com.appsmith.server.services.ConfigService;
 import com.appsmith.server.services.FeatureFlagService;
+import com.appsmith.server.services.OrganizationService;
 import com.appsmith.server.solutions.ReleaseNotesService;
 import com.appsmith.util.WebClientUtils;
 import joptsimple.internal.Strings;
@@ -54,8 +55,8 @@ public class InstanceConfigHelperCEImpl implements InstanceConfigHelperCE {
     private final AnalyticsService analyticsService;
     private final NetworkUtils networkUtils;
     private final ReleaseNotesService releaseNotesService;
-
     private final RTSCaller rtsCaller;
+    private final OrganizationService organizationService;
 
     private boolean isRtsAccessible = false;
 
@@ -225,8 +226,25 @@ public Mono<String> checkMongoDBVersion() {
                 });
     }
 
+    /**
+     * Method to trigger update for the organization feature flags. This method is called during the startup of
+     * the application. It's required at the startup to ensure that the feature flags are up-to-date which will then be
+     * consumed by {@link com.appsmith.server.aspect.FeatureFlaggedMethodInvokerAspect} in a non-reactive manner.
+     * In case the user tries to fetch the feature flags before the cache is updated, the aspect will fallback to the
+     * earlier cached data i.e. disabled state.
+     * @return  Empty Mono
+     */
     @Override
     public Mono<Void> updateCacheForOrganizationFeatureFlags() {
-        return featureFlagService.getOrganizationFeatures().then();
+        // TODO @CloudBilling: Fix this to update feature flags for all organizations and also should not affect the
+        //  startup
+        return organizationService
+                .retrieveAll()
+                .flatMap(org -> featureFlagService.getOrganizationFeatures(org.getId()))
+                .onErrorResume(error -> {
+                    log.error("Error while updating cache for org feature flags", error);
+                    return Mono.empty();
+                })
+                .then();
     }
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/migrations/JsonSchemaVersions.java

@@ -14,7 +14,7 @@
 public class JsonSchemaVersions extends JsonSchemaVersionsFallback {
 
     /**
-     * Only tenant level flags would work over here.
+     * Only org level flags would work over here.
      * @return an Integer which is server version
      */
     @Override

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/BaseRepository.java

@@ -2,6 +2,7 @@
 
 import org.springframework.data.mongodb.repository.ReactiveMongoRepository;
 import org.springframework.data.repository.NoRepositoryBean;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import java.io.Serializable;
@@ -34,4 +35,11 @@ public interface BaseRepository<T, ID extends Serializable> extends ReactiveMong
      * @return
      */
     Mono<Boolean> archiveAllById(Collection<ID> ids);
+
+    /**
+     * This function retrieves all the documents in the collection without the user context. This method will be used
+     * by internal server workflows like crons etc. to fetch all the documents from the collection.
+     * @return
+     */
+    Flux<T> retrieveAll();
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/BaseRepositoryImpl.java

@@ -101,6 +101,13 @@ public Flux<T> findAll() {
                 });
     }
 
+    @Override
+    public Flux<T> retrieveAll() {
+        Query query = new Query(notDeleted());
+        return mongoOperations.find(
+                query.cursorBatchSize(10000), entityInformation.getJavaType(), entityInformation.getCollectionName());
+    }
+
     /**
      * We don't use this today, it doesn't use our `notDeleted` criteria, and since we don't use it, we're not porting
      * it to Postgres. Querying with `queryBuilder` or anything else is arguably more readable than this.

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomOrganizationRepositoryCE.java

@@ -2,5 +2,8 @@
 
 import com.appsmith.server.domains.Organization;
 import com.appsmith.server.repositories.AppsmithRepository;
+import reactor.core.publisher.Mono;
 
-public interface CustomOrganizationRepositoryCE extends AppsmithRepository<Organization> {}
+public interface CustomOrganizationRepositoryCE extends AppsmithRepository<Organization> {
+    Mono<Integer> disableRestartForAllTenants();
+}

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomOrganizationRepositoryCEImpl.java

@@ -1,9 +1,22 @@
 package com.appsmith.server.repositories.ce;
 
 import com.appsmith.server.domains.Organization;
+import com.appsmith.server.helpers.ce.bridge.Bridge;
 import com.appsmith.server.repositories.BaseAppsmithRepositoryImpl;
 import lombok.extern.slf4j.Slf4j;
+import reactor.core.publisher.Mono;
+
+import static com.appsmith.server.domains.Organization.Fields.organizationConfiguration_isRestartRequired;
 
 @Slf4j
 public class CustomOrganizationRepositoryCEImpl extends BaseAppsmithRepositoryImpl<Organization>
-        implements CustomOrganizationRepositoryCE {}
+        implements CustomOrganizationRepositoryCE {
+
+    @Override
+    public Mono<Integer> disableRestartForAllTenants() {
+        log.info("Disabling restart for all tenants");
+        return queryBuilder()
+                .criteria(Bridge.isTrue(organizationConfiguration_isRestartRequired))
+                .updateAll(Bridge.update().set(organizationConfiguration_isRestartRequired, false));
+    }
+}

app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/FeatureFlagServiceCE.java

@@ -29,14 +29,17 @@ public interface FeatureFlagServiceCE {
      * To get all features of the organization from Cloud Services and store them locally
      * @return Mono of Void
      */
-    Mono<Void> getAllRemoteFeaturesForOrganizationAndUpdateFeatureFlagsWithPendingMigrations();
+    Mono<Organization> getAllRemoteFeaturesForOrganizationAndUpdateFeatureFlagsWithPendingMigrations(
+            Organization organization);
 
     /**
      * To get all features of the current organization.
      * @return Mono of Map
      */
     Mono<Map<String, Boolean>> getOrganizationFeatures();
 
+    Mono<Map<String, Boolean>> getOrganizationFeatures(String orgId);
+
     Mono<Organization> checkAndExecuteMigrationsForOrganizationFeatureFlags(Organization organization);
 
     CachedFeatures getCachedOrganizationFeatureFlags();