From 1b89420f996ac4b04bcdb49d71042e2cfdc6b76d Mon Sep 17 00:00:00 2001 From: Noah Levitt Date: Thu, 14 Mar 2019 13:47:15 -0700 Subject: [PATCH] trough dedup! --- contrib/pom.xml | 21 + .../postprocessor/TroughCrawlLogFeed.java | 61 ++- .../recrawl/TroughContentDigestHistory.java | 166 ++++++++ .../java/org/archive/trough/TroughClient.java | 375 ++++++++++++++++++ 4 files changed, 586 insertions(+), 37 deletions(-) create mode 100644 contrib/src/main/java/org/archive/modules/recrawl/TroughContentDigestHistory.java create mode 100644 contrib/src/main/java/org/archive/trough/TroughClient.java diff --git a/contrib/pom.xml b/contrib/pom.xml index 1a77567d..5a1e1ba6 100644 --- a/contrib/pom.xml +++ b/contrib/pom.xml @@ -61,6 +61,16 @@ kafka_2.10 0.9.0.0 + + com.rethinkdb + rethinkdb-driver + 2.3.3 + + + com.googlecode.json-simple + json-simple + 1.1.1 + @@ -105,6 +115,17 @@ + + + org.apache.maven.plugins + maven-compiler-plugin + 3.8.0 + + 1.8 + 1.8 + UTF-8 + + diff --git a/contrib/src/main/java/org/archive/modules/postprocessor/TroughCrawlLogFeed.java b/contrib/src/main/java/org/archive/modules/postprocessor/TroughCrawlLogFeed.java index 336e2d9f..492a9aa7 100644 --- a/contrib/src/main/java/org/archive/modules/postprocessor/TroughCrawlLogFeed.java +++ b/contrib/src/main/java/org/archive/modules/postprocessor/TroughCrawlLogFeed.java @@ -27,7 +27,6 @@ import java.util.logging.Logger; import org.apache.commons.collections.Closure; import org.apache.commons.io.IOUtils; -import org.apache.commons.lang.StringEscapeUtils; import org.apache.commons.lang.StringUtils; import org.apache.http.HttpResponse; import org.apache.http.client.methods.HttpPost; @@ -40,7 +39,7 @@ import org.archive.crawler.frontier.BdbFrontier; import org.archive.modules.CrawlURI; import org.archive.modules.Processor; import org.archive.modules.net.ServerCache; -import org.archive.util.ArchiveUtils; +import org.archive.trough.TroughClient; import org.archive.util.MimetypeUtils; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.Lifecycle; @@ -171,18 +170,6 @@ public class TroughCrawlLogFeed extends Processor implements Lifecycle { super.stop(); } - protected String sqlValue(Object o) { - if (o == null) { - return "null"; - } else if (o instanceof Date) { - return "datetime('" + ArchiveUtils.getLog14Date((Date) o) + "')"; - } else if (o instanceof Number) { - return o.toString(); - } else { - return "'" + StringEscapeUtils.escapeSql(o.toString()) + "'"; - } - } - transient protected CloseableHttpClient _httpClient; protected CloseableHttpClient httpClient() { if (_httpClient == null) { @@ -223,22 +210,22 @@ public class TroughCrawlLogFeed extends Processor implements Lifecycle { warcContentBytes = 0; } synchronized (crawledBatch) { - crawledBatch.add("(" + sqlValue(new Date(curi.getFetchBeginTime())) + ", " - + sqlValue(curi.getFetchStatus()) + ", " - + sqlValue(curi.getContentSize()) + ", " - + sqlValue(curi.getContentLength()) + ", " - + sqlValue(curi) + ", " - + sqlValue(curi.getPathFromSeed()) + ", " - + sqlValue(curi.isSeed() && !"".equals(curi.getPathFromSeed()) ? 1 : 0) + ", " - + sqlValue(curi.getVia()) + ", " - + sqlValue(MimetypeUtils.truncate(curi.getContentType())) + ", " - + sqlValue(curi.getContentDigestSchemeString()) + ", " - + sqlValue(curi.getSourceTag()) + ", " - + sqlValue(curi.isRevisit() ? 1 : 0) + ", " - + sqlValue(curi.getExtraInfo().opt("warcFilename")) + ", " - + sqlValue(curi.getExtraInfo().opt("warcOffset")) + ", " - + sqlValue(warcContentBytes) + ", " - + sqlValue(serverCache.getHostFor(curi.getUURI()).getHostName()) + ")"); + crawledBatch.add("(" + TroughClient.sqlValue(new Date(curi.getFetchBeginTime())) + ", " + + TroughClient.sqlValue((Object) curi.getFetchStatus()) + ", " + + TroughClient.sqlValue((Object) curi.getContentSize()) + ", " + + TroughClient.sqlValue((Object) curi.getContentLength()) + ", " + + TroughClient.sqlValue(curi) + ", " + + TroughClient.sqlValue(curi.getPathFromSeed()) + ", " + + TroughClient.sqlValue((Object) (curi.isSeed() && !"".equals(curi.getPathFromSeed()) ? 1 : 0)) + ", " + + TroughClient.sqlValue(curi.getVia()) + ", " + + TroughClient.sqlValue(MimetypeUtils.truncate(curi.getContentType())) + ", " + + TroughClient.sqlValue(curi.getContentDigestSchemeString()) + ", " + + TroughClient.sqlValue(curi.getSourceTag()) + ", " + + TroughClient.sqlValue((Object) (curi.isRevisit() ? 1 : 0)) + ", " + + TroughClient.sqlValue(curi.getExtraInfo().opt("warcFilename")) + ", " + + TroughClient.sqlValue(curi.getExtraInfo().opt("warcOffset")) + ", " + + TroughClient.sqlValue((Object) warcContentBytes) + ", " + + TroughClient.sqlValue(serverCache.getHostFor(curi.getUURI()).getHostName()) + ")"); if (crawledBatch.size() >= BATCH_MAX_SIZE || System.currentTimeMillis() - crawledBatchLastTime > BATCH_MAX_TIME_MS) { postCrawledBatch(); } @@ -246,13 +233,13 @@ public class TroughCrawlLogFeed extends Processor implements Lifecycle { } else { synchronized (uncrawledBatch) { uncrawledBatch.add("(" - + sqlValue(new Date()) + ", " - + sqlValue(curi) + ", " - + sqlValue(curi.getPathFromSeed()) + ", " - + sqlValue(curi.getFetchStatus()) + ", " - + sqlValue(curi.getVia()) + ", " - + sqlValue(curi.getSourceTag()) + ", " - + sqlValue(serverCache.getHostFor(curi.getUURI()).getHostName()) + ")"); + + TroughClient.sqlValue(new Date()) + ", " + + TroughClient.sqlValue(curi) + ", " + + TroughClient.sqlValue(curi.getPathFromSeed()) + ", " + + TroughClient.sqlValue((Object) curi.getFetchStatus()) + ", " + + TroughClient.sqlValue(curi.getVia()) + ", " + + TroughClient.sqlValue(curi.getSourceTag()) + ", " + + TroughClient.sqlValue(serverCache.getHostFor(curi.getUURI()).getHostName()) + ")"); if (uncrawledBatch.size() >= BATCH_MAX_SIZE || System.currentTimeMillis() - uncrawledBatchLastTime > BATCH_MAX_TIME_MS) { postUncrawledBatch(); } diff --git a/contrib/src/main/java/org/archive/modules/recrawl/TroughContentDigestHistory.java b/contrib/src/main/java/org/archive/modules/recrawl/TroughContentDigestHistory.java new file mode 100644 index 00000000..77acd192 --- /dev/null +++ b/contrib/src/main/java/org/archive/modules/recrawl/TroughContentDigestHistory.java @@ -0,0 +1,166 @@ +package org.archive.modules.recrawl; + +import static org.archive.modules.recrawl.RecrawlAttributeConstants.A_ORIGINAL_DATE; +import static org.archive.modules.recrawl.RecrawlAttributeConstants.A_ORIGINAL_URL; +import static org.archive.modules.recrawl.RecrawlAttributeConstants.A_WARC_RECORD_ID; + +import java.net.MalformedURLException; +import java.text.SimpleDateFormat; +import java.util.Date; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.TimeZone; +import java.util.logging.Level; +import java.util.logging.Logger; + +import org.archive.modules.CrawlURI; +import org.archive.modules.writer.WARCWriterProcessor; +import org.archive.spring.HasKeyedProperties; +import org.archive.spring.KeyedProperties; +import org.archive.trough.TroughClient; +import org.archive.util.ArchiveUtils; +import org.springframework.context.Lifecycle; + +/** + * AbstractContentDigestHistory implementation for trough. + * + *

To use, define a {@code TroughContentDigestHistory} top-level bean in your + * crawler-beans.cxml, then add {@link ContentDigestHistoryLoader} and + * {@link ContentDigestHistoryStorer} to your fetch chain, sandwiching the + * {@link WARCWriterProcessor}. In other words, follow the directions at + * https://github.com/internetarchive/heritrix3/wiki/Duplication%20Reduction%20Processors + * but replace the {@link BdbContentDigestHistory} bean with a + * {@code TroughContentDigestHistory} bean. + * + *

To understand how to use trough as a client, see: + *

+ * + * @see trough dedup implementation in warcprox + */ +public class TroughContentDigestHistory extends AbstractContentDigestHistory implements HasKeyedProperties, Lifecycle { + private static final Logger logger = Logger.getLogger(TroughContentDigestHistory.class.getName()); + + protected KeyedProperties kp = new KeyedProperties(); + public KeyedProperties getKeyedProperties() { + return kp; + } + + public void setSegmentId(String segmentId) { + kp.put("segmentId", segmentId); + } + public String getSegmentId() { + return (String) kp.get("segmentId"); + } + + /** + * @param rethinkUrl url with schema rethinkdb:// pointing to + * trough configuration database + */ + public void setRethinkUrl(String rethinkUrl) { + kp.put("rethinkUrl", rethinkUrl); + } + public String getRethinkUrl() { + return (String) kp.get("rethinkUrl"); + } + + protected TroughClient _troughClient = null; + protected boolean started = false; + + protected TroughClient troughClient() throws MalformedURLException { + if (_troughClient == null) { + _troughClient = new TroughClient(getRethinkUrl(), 60 * 60); + } + return _troughClient; + } + + protected static final String SCHEMA_ID = "warcprox-dedup-v1"; + protected static final String SCHEMA_SQL = "create table dedup (\n" + + " digest_key varchar(100) primary key,\n" + + " url varchar(2100) not null,\n" + + " date datetime not null,\n" + + " id varchar(100));\n"; // warc record id + + @Override + public void start() { + try { + troughClient().registerSchema(SCHEMA_ID, SCHEMA_SQL); + } catch (Exception e) { + // can happen. hopefully someone else has registered it + logger.log(Level.SEVERE, "will try to continue after problem registering schema " + SCHEMA_ID, e); + } + started = true; + } + + @Override + public void stop() { + } + + @Override + public boolean isRunning() { + return started; + } + + // dates come back from sqlite in this format: 2019-03-14 00:49:14 + protected static ThreadLocal SQLITE_DATE_FORMAT = new ThreadLocal() { + protected SimpleDateFormat initialValue() { + SimpleDateFormat df = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss", Locale.ENGLISH); + df.setTimeZone(TimeZone.getTimeZone("GMT")); + return df; + } + }; + + @Override + public void load(CrawlURI curi) { + // make this call in all cases so that the value is initialized and + // WARCWriterProcessor knows it should put the info in there + HashMap contentDigestHistory = curi.getContentDigestHistory(); + + try { + String sql = "select * from dedup where digest_key = %s"; + List> results = troughClient().read(getSegmentId(), sql, new String[] {persistKeyFor(curi)}); + if (!results.isEmpty()) { + Map hist = new HashMap(); + hist.put(A_ORIGINAL_URL, results.get(0).get("url")); + Date date = SQLITE_DATE_FORMAT.get().parse((String) results.get(0).get("date")); + hist.put(A_ORIGINAL_DATE, ArchiveUtils.getLog14Date(date)); + hist.put(A_WARC_RECORD_ID, results.get(0).get("id")); + + if (logger.isLoggable(Level.FINER)) { + logger.finer("loaded history by digest " + persistKeyFor(curi) + + " for uri " + curi + " - " + hist); + } + contentDigestHistory.putAll(hist); + } + } catch (Exception e) { + logger.log(Level.WARNING, "problem retrieving dedup info from trough segment " + getSegmentId() + " for url " + curi, e); + } + } + + protected static final String WRITE_SQL_TMPL = + "insert or ignore into dedup (digest_key, url, date, id) values (%s, %s, %s, %s);"; + + @Override + public void store(CrawlURI curi) { + if (!curi.hasContentDigestHistory() || curi.getContentDigestHistory().isEmpty()) { + return; + } + Map hist = curi.getContentDigestHistory(); + + try { + String digestKey = persistKeyFor(curi); + Object url = hist.get(A_ORIGINAL_URL); + Date date = ArchiveUtils.parse14DigitISODate((String) hist.get(A_ORIGINAL_DATE), null); + Object recordId = hist.get(A_WARC_RECORD_ID); + Object[] values = new Object[] { digestKey, url, date, recordId }; + troughClient().write(getSegmentId(), WRITE_SQL_TMPL, values, SCHEMA_ID); + } catch (Exception e) { + logger.log(Level.WARNING, "problem writing dedup info to trough segment " + getSegmentId() + " for url " + curi, e); + + } + } +} diff --git a/contrib/src/main/java/org/archive/trough/TroughClient.java b/contrib/src/main/java/org/archive/trough/TroughClient.java new file mode 100644 index 00000000..87979dd6 --- /dev/null +++ b/contrib/src/main/java/org/archive/trough/TroughClient.java @@ -0,0 +1,375 @@ +package org.archive.trough; + +import java.io.IOException; +import java.io.InputStreamReader; +import java.net.HttpURLConnection; +import java.net.MalformedURLException; +import java.net.URL; +import java.util.Date; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.Set; +import java.util.logging.Level; +import java.util.logging.Logger; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +import org.archive.util.ArchiveUtils; +import org.json.simple.JSONObject; +import org.json.simple.parser.JSONParser; +import org.json.simple.parser.ParseException; + +import com.rethinkdb.RethinkDB; +import com.rethinkdb.gen.ast.ReqlExpr; +import com.rethinkdb.net.Connection; + +public class TroughClient { + public static String sqlValue(Object o) { + if (o == null) { + return "null"; + } else if (o instanceof Date) { + return "datetime('" + ArchiveUtils.getLog14Date((Date) o) + "')"; + } else if (o instanceof Boolean) { + if ((Boolean) o) { + return "1"; + } else { + return "0"; + } + } else if (o instanceof Number) { + return o.toString(); + } else { + // the only character that needs escaped in sqlite string literals + // is single-quote, which is escaped as two single-quotes + return "'" + o.toString().replaceAll("'", "''") + "'"; + } + } + + private static final Logger logger = Logger.getLogger(TroughClient.class.getName()); + + protected static final RethinkDB r = RethinkDB.r; + + protected static final int SIX_HOURS_MS = 6 * 60 * 60 * 1000; + protected static final int TEN_MINUTES_MS = 10 * 60 * 1000; + protected static final String JSON_MIMETYPE = "application/json"; + protected static final String SQL_MIMETYPE = "application/sql"; + + protected Random rand = new Random(); + protected Map writeUrlCache; + protected Map readUrlCache; + protected Set dirtySegments; + protected String[] rethinkServers; + protected String rethinkDb; + protected Integer promotionInterval; + + public class TroughException extends IOException { + private static final long serialVersionUID = 1L; + public TroughException(String msg) { + super(msg); + } + public TroughException(Exception e) { + super(e); + } + public TroughException(String msg, Throwable cause) { + super(msg, cause); + } + } + + class Promotrix implements Runnable { + @Override + public void run() { + while (true) { + try { + Thread.sleep(promotionInterval * 1000); + + String[] promoteThese; + synchronized (dirtySegments) { + promoteThese = dirtySegments.toArray(new String[0]); + dirtySegments.clear(); + } + + if (promoteThese.length > 0) { + logger.info("promoting " + promoteThese.length + " trough segments"); + } + for (String segmentId: promoteThese) { + try { + promote(segmentId); + } catch (Exception e) { + logger.log(Level.WARNING, "problem promoting segment " + segmentId, e); + } + } + } catch (Exception e) { + logger.log(Level.WARNING, "continuing after unexpected exception in promoter thread", e); + } + } + } + } + + public TroughClient(String rethinkdbTroughUrl) throws MalformedURLException { + this(rethinkdbTroughUrl, null); + } + + protected HttpURLConnection httpRequest(String method, String url, String contentType, String payload, int timeout) throws IOException { + HttpURLConnection connection = (HttpURLConnection) new URL(url).openConnection(); + + connection.setConnectTimeout(timeout); + connection.setReadTimeout(timeout); + + connection.setRequestMethod(method); + + if (contentType != null) { + connection.setRequestProperty("Content-Type", contentType); + } + if (payload != null) { + connection.setDoOutput(true); + byte[] payloadBytes = payload.getBytes("UTF-8"); + logger.fine(method + " " + url + " " + payloadBytes.length + " bytes " + contentType); + connection.setRequestProperty("Content-Length", Integer.toString(payloadBytes.length)); + connection.getOutputStream().write(payloadBytes); + connection.getOutputStream().flush(); + connection.getOutputStream().close(); + } else { + logger.fine(method + " " + url); + } + + return connection; + } + + @SuppressWarnings("unchecked") + public void promote(String segmentId) throws IOException, TroughException { + String url = segmentManagerUrl() + "/promote"; + JSONObject payload = new JSONObject(); + payload.put("segment", segmentId); + + HttpURLConnection connection = httpRequest("POST", url, JSON_MIMETYPE, payload.toString(), SIX_HOURS_MS); + if (connection.getResponseCode() != 200) { + throw new TroughException("received " + connection.getResponseCode() + ": " + connection.getContent() + + " in response to POST " + url + " with data " + payload.toString()); + } + } + + /** + * Run a rethinkdb query. {@code retries=0} means try once. Default is 9 + * retries (10 tries). + * + * @return result of {@code reql.run()} + * @throws TroughException + */ + protected Object rethinkQuery(ReqlExpr reql, Integer retries) throws TroughException { + logger.fine("querying rethinkdb: " + reql); + if (retries == null || retries < 0) { + retries = 9; + } + + Exception lastE = null; + for (int i = 0; i <= retries; i++) { + int whichServer = rand.nextInt(rethinkServers.length); + try { + String[] hostPort = rethinkServers[whichServer].split(":", 2); + String host = hostPort[0]; + int port = 28015; + if (hostPort.length == 2) { + port = Integer.valueOf(hostPort[1]); + } + Connection conn = r.connection().hostname(host).port(port).db(rethinkDb).connect(); + Object result = reql.run(conn); + return result; + } catch (Exception e) { + logger.warning("rethinkdb query failed (server=" + rethinkServers[whichServer] + "; " + i + " retries left)"); + lastE = e; + } + } + + throw new TroughException(lastE); + } + + protected String segmentManagerUrl() throws TroughException { + ReqlExpr reql = r.table("services").optArg("read_mode", "majority").getAll("trough-sync-master") + .filter(svc -> r.now().sub(svc.g("last_heartbeat")).lt(svc.g("ttl"))); + @SuppressWarnings("unchecked") + Iterable> results = (Iterable>) rethinkQuery(reql, null); + for (Map result: results) { + return (String) result.get("url"); + } + throw new TroughException("no healthy trough-sync-master in rethinkdb?"); + } + + /** + * + * @param rethinkdbTroughUrl url with schema rethinkdb:// pointing to + * trough configuration database + * @param promotionInterval if specified, {@code TroughClient} will spawn a + * thread that "promotes" (pushed to hdfs) "dirty" trough segments + * (segments that have received writes) periodically, sleeping for + * {@code promotionInterval} seconds between cycles + * @throws MalformedURLException + */ + public TroughClient(String rethinkdbTroughUrl, Integer promotionInterval) throws MalformedURLException { + parseRethinkdbUrl(rethinkdbTroughUrl); + writeUrlCache = new HashMap(); + readUrlCache = new HashMap(); + dirtySegments = new HashSet(); + + if (promotionInterval != null) { + this.promotionInterval = promotionInterval; + Thread promotrix = new Thread(new Promotrix(), "TroughClient-promotrix"); + promotrix.setDaemon(true); + promotrix.start(); + } + } + + /** + * Parses a url like this rethinkdb://server1:port,server2:port/database/table + * Sets fields {@code rethinkServers}, {@code rethinkDb}, {@code rethinkTable} + * @throws MalformedURLException + */ + protected void parseRethinkdbUrl(String input) throws MalformedURLException { + Matcher m = Pattern.compile("^rethinkdb://([^/]+)/([^/]+)$").matcher(input); + if (!m.matches()) { + throw new MalformedURLException("failed to parse as rethinkdb url: " + input); + } + + rethinkServers = m.group(1).split(","); + rethinkDb = m.group(2); + } + + // XXX inherits inconsistency in handling exceptions from warcprox trough client + @SuppressWarnings("unchecked") + public List> read(String segmentId, String sqlTmpl, Object[] values) throws IOException { + String readUrl = readUrl(segmentId); + + String[] sqlValues = new String[values.length]; + for (int i = 0; i < values.length; i++) { + sqlValues[i] = sqlValue(values[i]); + } + String sql = String.format(sqlTmpl, (Object[]) sqlValues); + + HttpURLConnection connection; + try { + connection = httpRequest("POST", readUrl, SQL_MIMETYPE, sql, TEN_MINUTES_MS); + if (connection.getResponseCode() != 200) { + readUrlCache.remove(segmentId); + throw new TroughException( + "unexpected response" + connection.getResponseCode() + " " + + connection.getResponseMessage() + ": " + connection.getContent() + + " from " + readUrl + " to query: " + sql); + } + + Object result = new JSONParser().parse(new InputStreamReader(connection.getInputStream(), "UTF-8")); + return (List>) result; + } catch (IOException e) { + readUrlCache.remove(segmentId); + throw e; + } catch (ParseException e) { + readUrlCache.remove(segmentId); + throw new TroughException("problem parsing json response from " + readUrl, e); + } + } + + protected String readUrl(String segmentId) throws TroughException { + if (readUrlCache.get(segmentId) == null) { + String url = readUrlNoCache(segmentId); + readUrlCache.put(segmentId, url); + logger.info("segment " + segmentId + " read url is " + url); + } + return readUrlCache.get(segmentId); + } + + @SuppressWarnings("unchecked") + protected String readUrlNoCache(String segmentId) throws TroughException { + ReqlExpr reql = r.table("services").getAll(segmentId).optArg("index", "segment") + .filter(r.hashMap("role", "trough-read")) + .filter(svc -> r.now().sub(svc.g("last_heartbeat")).lt(svc.g("ttl"))) + .orderBy("load"); + List> result; + result = (List>) rethinkQuery(reql, null); + if (result != null && result.size() > 0) { + return (String) result.get(0).get("url"); + } else { + throw new TroughException("failed to obtain read url for trough segment " + segmentId); + } + } + + public void registerSchema(String schemaId, String schemaSql) throws IOException { + String url = segmentManagerUrl() + "/schema/" + schemaId + "/sql"; + HttpURLConnection connection = httpRequest("PUT", url, SQL_MIMETYPE, schemaSql, TEN_MINUTES_MS); + if (connection.getResponseCode() != 201 && connection.getResponseCode() != 204) { + throw new TroughException("received " + connection.getResponseCode() + ": " + connection.getContent() + + " in response to PUT " + url + " with data " + schemaSql); + } + } + + public void write(String segmentId, String sqlTmpl, Object[] values) throws IOException { + write(segmentId, sqlTmpl, values, "default"); + } + + public void write(String segmentId, String sqlTmpl, Object[] values, String schemaId) throws IOException { + String url = writeUrl(segmentId, schemaId); + + String[] sqlValues = new String[values.length]; + for (int i = 0; i < values.length; i++) { + sqlValues[i] = sqlValue(values[i]); + } + String sql = String.format(sqlTmpl, (Object[]) sqlValues); + + try { + HttpURLConnection connection = httpRequest("POST", url, "application/sql", sql, TEN_MINUTES_MS); + if (connection.getResponseCode() != 200) { + throw new TroughException("unexpected response" + connection.getResponseCode() + " " + + connection.getResponseMessage() + ": " + connection.getContent() + + " from " + url + " to query: " + sql); + } + if (!dirtySegments.contains(segmentId)) { + synchronized (segmentId) { + dirtySegments.add(segmentId); + } + } + } catch (IOException e) { + writeUrlCache.remove("segmentId"); + throw e; + } + } + + protected String writeUrl(String segmentId, String schemaId) throws IOException { + if (writeUrlCache.get(segmentId) == null) { + String url = writeUrlNoCache(segmentId, schemaId); + writeUrlCache.put(segmentId, url); + logger.info("segment " + segmentId + " write url is " + url); + } + return writeUrlCache.get(segmentId); + } + + @SuppressWarnings("unchecked") + protected String writeUrlNoCache(String segmentId, String schemaId) throws IOException { + String provisionUrl; + String segmentManagerUrl = segmentManagerUrl(); + if (segmentManagerUrl.endsWith("/")) { + provisionUrl = segmentManagerUrl + "provision"; + } else { + provisionUrl = segmentManagerUrl + "/provision"; + } + + JSONObject payload = new JSONObject(); + payload.put("segment", segmentId); + payload.put("schema", schemaId); + HttpURLConnection connection = httpRequest("POST", provisionUrl, JSON_MIMETYPE, payload.toJSONString(), TEN_MINUTES_MS); + if (connection.getResponseCode() != 200) { + throw new TroughException("received " + connection.getResponseCode() + ": " + connection.getContent() + + " in response to POST " + provisionUrl + " with data " + payload); + } + + JSONObject result; + try { + result = (JSONObject) new JSONParser().parse(new InputStreamReader(connection.getInputStream(), "UTF-8")); + } catch (ParseException e) { + throw new TroughException("unable to parse response from POST " + provisionUrl + " as json", e); + } + String writeUrl = (String) result.get("write_url"); + if (writeUrl == null) { + throw new TroughException("write_url missing from response to " + provisionUrl + " - " + result); + } + return writeUrl; + } +}