@Override public void putAll(String storeName, String tableName, Collection<T> states) throws IOException { try (ByteArrayOutputStream byteArrayOs = new ByteArrayOutputStream(); OutputStream os = compressedValues ? new GZIPOutputStream(byteArrayOs) : byteArrayOs; DataOutputStream dataOutput = new DataOutputStream(os)) { for (T state : states) { addStateToDataOutputStream(dataOutput, state); } dataOutput.close(); putData(storeName, tableName, byteArrayOs.toByteArray()); } }
@Override public List<T> getAll(String storeName, String tableName) throws IOException { List<T> states = Lists.newArrayList(); String path = formPath(storeName, tableName); byte[] data = propStore.get(path, null, 0); deserialize(data, states); return states; }
@Override public List<T> getAll(String storeName) throws IOException { return getAll(storeName, Predicates.<String>alwaysTrue()); }
@Override public void createAlias(String storeName, String original, String alias) throws IOException { String pathOriginal = formPath(storeName, original); byte[] data; if (!propStore.exists(pathOriginal, 0)) { throw new IOException(String.format("State does not exist for table %s", original)); } data = propStore.get(pathOriginal, null, 0); putData(storeName, alias, data); }
@Override public void delete(String storeName, String tableName) throws IOException { propStore.remove(formPath(storeName, tableName), 0); }
@Override public <T extends State> StateStore<T> createStateStore(Config config, Class<T> stateClass) { String connectString = ConfigUtils.getString(config, ZkStateStoreConfigurationKeys.STATE_STORE_ZK_CONNECT_STRING_KEY, ZkStateStoreConfigurationKeys.STATE_STORE_ZK_CONNECT_STRING_DEFAULT); String rootDir = config.getString(ConfigurationKeys.STATE_STORE_ROOT_DIR_KEY); boolean compressedValues = ConfigUtils.getBoolean(config, ConfigurationKeys.STATE_STORE_COMPRESSED_VALUES_KEY, ConfigurationKeys.DEFAULT_STATE_STORE_COMPRESSED_VALUES); try { return new ZkStateStore(connectString, rootDir, compressedValues, stateClass); } catch (Exception e) { throw new RuntimeException("Failed to create ZkStateStore with factory", e); } } }
/** * Deserialize data into a list of {@link State}s. * @param data byte array * @param states output list of states * @throws IOException */ private void deserialize(byte[] data, List<T> states) throws IOException { deserialize(data, states, null); } }
@Override public void put(String storeName, String tableName, T state) throws IOException { putAll(storeName, tableName, Collections.singletonList(state)); }
@Override public boolean exists(String storeName, String tableName) throws IOException { String path = formPath(storeName, tableName); return propStore.exists(path, 0); }
@BeforeClass public void setUp() throws Exception { ConfigBuilder configBuilder = ConfigBuilder.create(); testingServer = new TestingServer(-1); zkJobStateStore = new ZkStateStore<>(testingServer.getConnectString(), "/STATE_STORE/TEST", false, JobState.class); configBuilder.addPrimitive(ZkStateStoreConfigurationKeys.STATE_STORE_ZK_CONNECT_STRING_KEY, testingServer.getConnectString()); configBuilder.addPrimitive(ConfigurationKeys.STATE_STORE_ROOT_DIR_KEY, "/STATE_STORE/TEST2"); ClassAliasResolver<DatasetStateStore.Factory> resolver = new ClassAliasResolver<>(DatasetStateStore.Factory.class); DatasetStateStore.Factory stateStoreFactory = resolver.resolveClass("zk").newInstance(); zkDatasetStateStore = stateStoreFactory.createStateStore(configBuilder.build()); // clear data that may have been left behind by a prior test run zkJobStateStore.delete(TEST_JOB_NAME); zkDatasetStateStore.delete(TEST_JOB_NAME); zkJobStateStore.delete(TEST_JOB_NAME2); zkDatasetStateStore.delete(TEST_JOB_NAME2); }
@Override public T get(String storeName, String tableName, String stateId) throws IOException { String path = formPath(storeName, tableName); byte[] data = propStore.get(path, null, 0); List<T> states = Lists.newArrayList(); deserialize(data, states, stateId); if (states.isEmpty()) { return null; } else { return states.get(0); } }
@Override public void delete(String storeName) throws IOException { propStore.remove(formPath(storeName), 0); }
/** * Retrieve states from the state store based on the store name and a filtering predicate * @param storeName The store name enclosing the state files * @param predicate The predicate for state file filtering * @return list of states matching matching the predicate * @throws IOException */ protected List<T> getAll(String storeName, Predicate<String> predicate) throws IOException { List<T> states = Lists.newArrayList(); String path = formPath(storeName); byte[] data; List<String> children = propStore.getChildNames(path, 0); if (children == null) { return Collections.emptyList(); } for (String c : children) { if (predicate.apply(c)) { data = propStore.get(path + "/" + c, null, 0); deserialize(data, states); } } return states; }
@Override public List<String> getTableNames(String storeName, Predicate<String> predicate) throws IOException { List<String> names = Lists.newArrayList(); String path = formPath(storeName); List<String> children = propStore.getChildNames(path, 0); if (children != null) { for (String c : children) { if (predicate.apply(c)) { names.add(c); } } } return names; }
@Override public boolean create(String storeName, String tableName) throws IOException { String path = formPath(storeName, tableName); if (propStore.exists(path, 0)) { throw new IOException(String.format("State already exists for storeName %s tableName %s", storeName, tableName)); } return propStore.create(path, ArrayUtils.EMPTY_BYTE_ARRAY, AccessOption.PERSISTENT); }
/** * Get store names in the state store * * @param predicate only returns names matching predicate * @return (possibly empty) list of store names from the given store * @throws IOException */ public List<String> getStoreNames(Predicate<String> predicate) throws IOException { List<String> names = Lists.newArrayList(); String path = formPath(""); List<String> children = propStore.getChildNames(path, 0); if (children != null) { for (String c : children) { if (predicate.apply(c)) { names.add(c); } } } return names; }
@Override public boolean create(String storeName) throws IOException { String path = formPath(storeName); return propStore.exists(path, 0) || propStore.create(path, ArrayUtils.EMPTY_BYTE_ARRAY, AccessOption.PERSISTENT); }
/** * Create a new znode with data if it does not exist otherwise update with data * @param storeName storeName portion of znode path * @param tableName tableName portion of znode path * @param data znode data * @throws IOException */ private void putData(String storeName, String tableName, byte[] data) throws IOException { String path = formPath(storeName, tableName); if (!propStore.exists(path, 0)) { // create with data if (!propStore.create(path, data, AccessOption.PERSISTENT)) { throw new IOException("Failed to create a state file for table " + tableName); } } else { // Update propStore.set(path, data, AccessOption.PERSISTENT); } }