<commit_before>package com.neverwinterdp.scribengin.dataflow;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import com.neverwinterdp.scribengin.dataflow.worker.DataflowTaskExecutorService;
import com.neverwinterdp.scribengin.storage.StreamDescriptor;
import com.neverwinterdp.scribengin.storage.sink.Sink;
import com.neverwinterdp.scribengin.storage.sink.SinkFactory;
import com.neverwinterdp.scribengin.storage.sink.SinkStream;
import com.neverwinterdp.scribengin.storage.sink.SinkStreamWriter;
import com.neverwinterdp.scribengin.storage.source.Source;
import com.neverwinterdp.scribengin.storage.source.SourceFactory;
import com.neverwinterdp.scribengin.storage.source.SourceStream;
import com.neverwinterdp.scribengin.storage.source.SourceStreamReader;
import com.neverwinterdp.yara.Meter;

public class DataflowTaskContext {
  private DataflowTaskExecutorService executorService;
  private DataflowTaskReport report;
  private SourceContext sourceContext;
  private Map<String, SinkContext> sinkContexts = new HashMap<String, SinkContext>();
  private DataflowTaskDescriptor dataflowTaskDescriptor;
  private boolean complete = false;
  private Meter dataflowReadMeter ;
  
  public DataflowTaskContext(DataflowTaskExecutorService service, DataflowTaskDescriptor descriptor, DataflowTaskReport report)  throws Exception {
    executorService = service;
    sourceContext = new SourceContext(service.getSourceFactory(), descriptor.getSourceStreamDescriptor());
    dataflowReadMeter = service.getMetricRegistry().getMeter("dataflow.source.throughput", "byte") ;
    Iterator<Map.Entry<String, StreamDescriptor>> i = descriptor.getSinkStreamDescriptors().entrySet().iterator();
      while (i.hasNext()) {
        Map.Entry<String, StreamDescriptor> entry = i.next();
        StreamDescriptor streamDescriptor = entry.getValue();
        SinkContext context = new SinkContext(service.getSinkFactory(), streamDescriptor);
        sinkContexts.put(entry.getKey(), context);
      }
      this.report = report;
      this.dataflowTaskDescriptor = descriptor;
  }

  public DataflowTaskReport getReport() { return this.report; }
  
  public SourceStreamReader getSourceStreamReader() { return sourceContext.assignedSourceStreamReader; }

  public boolean isComplete() { return this.complete ; }
  
  public void setComplete(boolean b) { this.complete = b ; }
  
  public boolean isEndOfDataStream() throws Exception { 
    return  getSourceStreamReader().isEndOfDataStream() ; 
  }
  
  public DataflowMessage nextRecord(long maxWaitForDataRead) throws Exception {
    DataflowMessage dataflowMessage = getSourceStreamReader().next(maxWaitForDataRead);
    if(dataflowMessage != null) {
      dataflowReadMeter.mark(dataflowMessage.getData().length + dataflowMessage.getKey().length());
    }
    return dataflowMessage ;
  }
  
  public void append(DataflowMessage dataflowMessage) throws Exception {
    SinkContext sinkContext = sinkContexts.get("default");
    sinkContext.assignedSinkStreamWriter.append(dataflowMessage);
    Meter meter = executorService.getMetricRegistry().getMeter("dataflow.sink.default.throughput"  , "byte") ;
    meter.mark(dataflowMessage.getData().length + dataflowMessage.getKey().length());
  }

  public void write(String sinkName, DataflowMessage dataflowMessage) throws Exception {
    SinkContext sinkContext = sinkContexts.get(sinkName);
    sinkContext.assignedSinkStreamWriter.append(dataflowMessage);
    Meter meter = executorService.getMetricRegistry().getMeter("dataflow.sink." + sinkName + ".throughput"  , "byte") ;
    meter.mark(dataflowMessage.getData().length + dataflowMessage.getKey().length());
  }
  
  public String[] getAvailableSinks() {
    String[] name = new String[sinkContexts.size()] ;
    sinkContexts.keySet().toArray(name) ;
    return name ;
  }
  
  public SinkContext getSinkContext(String name) {
    return sinkContexts.get(name) ;
  }
  
  private void prepareCommit() throws Exception {
    sourceContext.assignedSourceStreamReader.prepareCommit();
    Iterator<SinkContext> i = sinkContexts.values().iterator();
    while (i.hasNext()) {
      SinkContext ctx = i.next();
      ctx.prepareCommit();
    }
  }

  public boolean commit() throws Exception {
    //prepareCommit is a vote to make sure both sink, invalidSink, and source
    //are ready to commit data, otherwise rollback will occur
    try {
      prepareCommit();
      completeCommit();
      report.updateCommit();
      return true;
    } catch (Exception ex) {
      report.setCommitFailCount(report.getCommitFailCount() + 1);
      rollback();
      throw ex;
    } finally {
      executorService.getDataflowRegistry().dataflowTaskReport(dataflowTaskDescriptor, report);
    }
  }

  private void rollback() throws Exception {
    //TODO: implement the proper transaction
    Iterator<SinkContext> i = sinkContexts.values().iterator();
    while (i.hasNext()) {
      SinkContext ctx = i.next();
      ctx.rollback();
    }
    sourceContext.rollback();
  }

  public void close() throws Exception {<commit_msg>fix commit and rollback<commit_after>package com.neverwinterdp.scribengin.dataflow;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import com.neverwinterdp.scribengin.dataflow.worker.DataflowTaskExecutorService;
import com.neverwinterdp.scribengin.storage.StreamDescriptor;
import com.neverwinterdp.scribengin.storage.sink.Sink;
import com.neverwinterdp.scribengin.storage.sink.SinkFactory;
import com.neverwinterdp.scribengin.storage.sink.SinkStream;
import com.neverwinterdp.scribengin.storage.sink.SinkStreamWriter;
import com.neverwinterdp.scribengin.storage.source.Source;
import com.neverwinterdp.scribengin.storage.source.SourceFactory;
import com.neverwinterdp.scribengin.storage.source.SourceStream;
import com.neverwinterdp.scribengin.storage.source.SourceStreamReader;
import com.neverwinterdp.yara.Meter;

public class DataflowTaskContext {
  private DataflowTaskExecutorService executorService;
  private DataflowTaskReport report;
  private SourceContext sourceContext;
  private Map<String, SinkContext> sinkContexts = new HashMap<String, SinkContext>();
  private DataflowTaskDescriptor dataflowTaskDescriptor;
  private boolean complete = false;
  private Meter dataflowReadMeter ;
  
  public DataflowTaskContext(DataflowTaskExecutorService service, DataflowTaskDescriptor descriptor, DataflowTaskReport report)  throws Exception {
    executorService = service;
    sourceContext = new SourceContext(service.getSourceFactory(), descriptor.getSourceStreamDescriptor());
    dataflowReadMeter = service.getMetricRegistry().getMeter("dataflow.source.throughput", "byte") ;
    Iterator<Map.Entry<String, StreamDescriptor>> i = descriptor.getSinkStreamDescriptors().entrySet().iterator();
      while (i.hasNext()) {
        Map.Entry<String, StreamDescriptor> entry = i.next();
        StreamDescriptor streamDescriptor = entry.getValue();
        SinkContext context = new SinkContext(service.getSinkFactory(), streamDescriptor);
        sinkContexts.put(entry.getKey(), context);
      }
      this.report = report;
      this.dataflowTaskDescriptor = descriptor;
  }

  public DataflowTaskReport getReport() { return this.report; }
  
  public SourceStreamReader getSourceStreamReader() { return sourceContext.assignedSourceStreamReader; }

  public boolean isComplete() { return this.complete ; }
  
  public void setComplete(boolean b) { this.complete = b ; }
  
  public boolean isEndOfDataStream() throws Exception { 
    return  getSourceStreamReader().isEndOfDataStream() ; 
  }
  
  public DataflowMessage nextRecord(long maxWaitForDataRead) throws Exception {
    DataflowMessage dataflowMessage = getSourceStreamReader().next(maxWaitForDataRead);
    if(dataflowMessage != null) {
      dataflowReadMeter.mark(dataflowMessage.getData().length + dataflowMessage.getKey().length());
    }
    return dataflowMessage ;
  }
  
  public void append(DataflowMessage dataflowMessage) throws Exception {
    SinkContext sinkContext = sinkContexts.get("default");
    sinkContext.assignedSinkStreamWriter.append(dataflowMessage);
    Meter meter = executorService.getMetricRegistry().getMeter("dataflow.sink.default.throughput"  , "byte") ;
    meter.mark(dataflowMessage.getData().length + dataflowMessage.getKey().length());
  }

  public void write(String sinkName, DataflowMessage dataflowMessage) throws Exception {
    SinkContext sinkContext = sinkContexts.get(sinkName);
    sinkContext.assignedSinkStreamWriter.append(dataflowMessage);
    Meter meter = executorService.getMetricRegistry().getMeter("dataflow.sink." + sinkName + ".throughput"  , "byte") ;
    meter.mark(dataflowMessage.getData().length + dataflowMessage.getKey().length());
  }
  
  public String[] getAvailableSinks() {
    String[] name = new String[sinkContexts.size()] ;
    sinkContexts.keySet().toArray(name) ;
    return name ;
  }
  
  public SinkContext getSinkContext(String name) {
    return sinkContexts.get(name) ;
  }
  
  private void prepareCommit() throws Exception {
    sourceContext.assignedSourceStreamReader.prepareCommit();
    Iterator<SinkContext> i = sinkContexts.values().iterator();
    while (i.hasNext()) {
      SinkContext ctx = i.next();
      ctx.prepareCommit();
    }
  }

  public boolean commit() throws Exception {
    //prepareCommit is a vote to make sure both sink, invalidSink, and source
    //are ready to commit data, otherwise rollback will occur
    try {
      prepareCommit();
      completeCommit();
      report.updateCommit();
      return true ;
    } catch (Exception ex) {
      report.setCommitFailCount(report.getCommitFailCount() + 1);
      rollback();
      return false ;
    } finally {
      executorService.getDataflowRegistry().dataflowTaskReport(dataflowTaskDescriptor, report);
    }
  }

  private void rollback() throws Exception {
    //TODO: implement the proper transaction
    Iterator<SinkContext> i = sinkContexts.values().iterator();
    while (i.hasNext()) {
      SinkContext ctx = i.next();
      ctx.rollback();
    }
    sourceContext.rollback();
  }

  public void close() throws Exception {
{'fraction_non_alphanumeric': 0.10053013940702926, 'fraction_numerical': 0.00019634792852935403, 'mean_word_length': 4.678372352285396, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2774856', 'n_tokens_mistral': 2827, 'n_tokens_neox': 2874, 'n_words': 561}