Map<Integer, Long> shardsForCounter(String counterId) throws IOException { final List<Key> shardKeys = IntStream.range(0, NUM_SHARDS).mapToObj( index -> datastore.newKeyFactory().setKind(KIND_COUNTER_SHARD).newKey( String.format("%s-%d", counterId, index))) .collect(toList()); final Map<Integer, Long> fetchedShards = new HashMap<>(); datastore.get(shardKeys, shard -> fetchedShards.put( (int) shard.getLong(PROPERTY_SHARD_INDEX), shard.getLong(PROPERTY_SHARD_VALUE))); return fetchedShards; }
Optional<Workflow> workflow(WorkflowId workflowId) throws IOException { final Optional<Entity> entityOptional = getOpt(datastore, workflowKey(datastore.newKeyFactory(), workflowId)) .filter(e -> e.contains(PROPERTY_WORKFLOW_JSON)); if (entityOptional.isPresent()) { return Optional.of(parseWorkflowJson(entityOptional.get(), workflowId)); } else { return Optional.empty(); } }
boolean enabled(WorkflowId workflowId) throws IOException { final Key workflowKey = workflowKey(datastore.newKeyFactory(), workflowId); return getOpt(datastore, workflowKey) .filter(w -> w.contains(PROPERTY_WORKFLOW_ENABLED)) .map(workflow -> workflow.getBoolean(PROPERTY_WORKFLOW_ENABLED)) .orElse(DEFAULT_WORKFLOW_ENABLED); }
@Override public Optional<Workflow> workflow(WorkflowId workflowId) throws IOException { final Optional<Entity> entityOptional = DatastoreStorage.getOpt(tx, DatastoreStorage.workflowKey(tx.getDatastore().newKeyFactory(), workflowId)); if (entityOptional.isPresent()) { return Optional.of(DatastoreStorage.parseWorkflowJson(entityOptional.get(), workflowId)); } else { return Optional.empty(); } }
@Override public Optional<Backfill> backfill(String id) throws IOException { final Key key = DatastoreStorage.backfillKey(tx.getDatastore().newKeyFactory(), id); final Entity entity = tx.get(key); if (entity == null) { return Optional.empty(); } return Optional.of(entityToBackfill(entity)); } }
Optional<Resource> getResource(String id) throws IOException { Entity entity = datastore.get(datastore.newKeyFactory().setKind(KIND_COUNTER_LIMIT).newKey(id)); if (entity == null) { return Optional.empty(); } return Optional.of(entityToResource(entity)); }
@Override public Optional<RunState> readActiveState(WorkflowInstance instance) throws IOException { final Entity entity = tx.get(activeWorkflowInstanceKey(tx.getDatastore().newKeyFactory(), instance)); if (entity == null) { return Optional.empty(); } else { return Optional.of(entityToRunState(entity, instance)); } }
Optional<Backfill> getBackfill(String id) throws IOException { final Entity entity = datastore.get(datastore.newKeyFactory().setKind(KIND_BACKFILL).newKey(id)); if (entity == null) { return Optional.empty(); } return Optional.of(entityToBackfill(entity)); }
@Override public WorkflowInstance updateActiveState(WorkflowInstance instance, RunState state) throws IOException { tx.update(runStateToEntity(tx.getDatastore().newKeyFactory(), instance, state)); return instance; }
StyxConfig config() throws IOException { final Entity entity = asBuilderOrNew( getOpt(datastore, globalConfigKey(datastore.newKeyFactory())), globalConfigKey(datastore.newKeyFactory())) .build(); return entityToConfig(entity); }
void delete(WorkflowId workflowId) throws IOException { storeWithRetries(() -> { datastore.delete(workflowKey(datastore.newKeyFactory(), workflowId)); return null; }); }
@Override public WorkflowInstance deleteActiveState(WorkflowInstance instance) throws IOException { tx.delete(activeWorkflowInstanceIndexShardEntryKey(tx.getDatastore().newKeyFactory(), instance)); tx.delete(activeWorkflowInstanceKey(tx.getDatastore().newKeyFactory(), instance)); return instance; }
long getLimitForCounter(String counterId) throws IOException { if (GLOBAL_RESOURCE_ID.equals(counterId)) { // missing global resource means free to go return config().globalConcurrency().orElse(Long.MAX_VALUE); } final Key limitKey = datastore.newKeyFactory().setKind(KIND_COUNTER_LIMIT).newKey(counterId); final Entity limitEntity = datastore.get(limitKey); if (limitEntity == null) { throw new IllegalArgumentException("No limit found in Datastore for " + counterId); } else { return limitEntity.getLong(PROPERTY_LIMIT); } }
public WorkflowState workflowState(WorkflowId workflowId) throws IOException { final WorkflowState.Builder builder = WorkflowState.builder(); final Optional<Entity> workflowEntity = getOpt(datastore, workflowKey(datastore.newKeyFactory(), workflowId)); builder.enabled(workflowEntity.filter(w -> w.contains(PROPERTY_WORKFLOW_ENABLED)) .map(workflow -> workflow.getBoolean(PROPERTY_WORKFLOW_ENABLED)) .orElse(DEFAULT_WORKFLOW_ENABLED)); getOptInstantProperty(workflowEntity, PROPERTY_NEXT_NATURAL_TRIGGER) .ifPresent(builder::nextNaturalTrigger); getOptInstantProperty(workflowEntity, PROPERTY_NEXT_NATURAL_OFFSET_TRIGGER) .ifPresent(builder::nextNaturalOffsetTrigger); return builder.build(); }
@Override public WorkflowInstance writeActiveState(WorkflowInstance instance, RunState state) throws IOException { // Note: the parent entity need not actually exist final Key indexEntryKey = activeWorkflowInstanceIndexShardEntryKey(tx.getDatastore().newKeyFactory(), instance); final Entity indexEntry = Entity.newBuilder(indexEntryKey).build(); tx.add(indexEntry); tx.add(runStateToEntity(tx.getDatastore().newKeyFactory(), instance, state)); return instance; }
@Override public Optional<Shard> shard(String counterId, int shardIndex) throws IOException { // TODO there's no need for this to be transactional final Key shardKey = tx.getDatastore().newKeyFactory().setKind(KIND_COUNTER_SHARD) .newKey(counterId + "-" + shardIndex); Entity shardEntity = tx.get(shardKey); if (shardEntity == null) { return Optional.empty(); } return Optional.of(Shard.create(counterId, shardIndex, (int) tx.get(shardKey).getLong(PROPERTY_SHARD_VALUE))); }
/** * Delete resource by id. Deletes both counter shards and counter limit if it exists. * * <p>Due to Datastore limitations (modify max 25 entity groups per transaction), * we cannot do everything in one transaction. */ void deleteResource(String id) throws IOException { storeWithRetries(() -> { datastore.delete(datastore.newKeyFactory().setKind(KIND_COUNTER_LIMIT).newKey(id)); return null; }); deleteShardsForCounter(id); }
@Override public void updateLimitForCounter(String counterId, long limit) throws IOException { final Key limitKey = tx.getDatastore().newKeyFactory().setKind(KIND_COUNTER_LIMIT).newKey(counterId); tx.put(Entity.newBuilder(limitKey).set(PROPERTY_LIMIT, limit).build()); }
private Entity resourceToEntity(CheckedDatastore datastore, Resource resource) { final Key key = datastore.newKeyFactory().setKind(KIND_COUNTER_LIMIT).newKey(resource.id()); return Entity.newBuilder(key) .set(PROPERTY_LIMIT, resource.concurrency()) .build(); }
@Override public void store(Shard shard) throws IOException { tx.put(Entity.newBuilder(tx.getDatastore().newKeyFactory().setKind(KIND_COUNTER_SHARD) .newKey(shard.counterId() + "-" + shard.index())) .set(PROPERTY_COUNTER_ID, shard.counterId()) .set(PROPERTY_SHARD_INDEX, shard.index()) .set(PROPERTY_SHARD_VALUE, shard.value()) .build()); }