/*
 * Copyright 2009-2010 by The Regents of the University of California
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * you may obtain a copy of the License from
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package edu.uci.ics.hyracks.control.cc.dataset;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

import edu.uci.ics.hyracks.api.comm.NetworkAddress;
import edu.uci.ics.hyracks.api.dataset.DatasetDirectoryRecord;
import edu.uci.ics.hyracks.api.dataset.IDatasetDirectoryService;
import edu.uci.ics.hyracks.api.dataset.ResultSetId;
import edu.uci.ics.hyracks.api.exceptions.HyracksDataException;
import edu.uci.ics.hyracks.api.job.JobId;

/**
 * TODO(madhusudancs): The potential perils of this global dataset directory service implementation is that, the jobs
 * location information is never evicted from the memory and the memory usage grows as the number of jobs in the system
 * grows. What we should possibly do is, add an API call for the client to say that it received everything it has to for
 * the job (after it receives all the results) completely. Then we can just get rid of the location information for that
 * job.
 */
public class DatasetDirectoryService implements IDatasetDirectoryService {
    private final Map<JobId, Map<ResultSetId, DatasetDirectoryRecord[]>> jobResultLocationsMap;

    public DatasetDirectoryService() {
        jobResultLocationsMap = new HashMap<JobId, Map<ResultSetId, DatasetDirectoryRecord[]>>();
    }

    @Override
    public synchronized void registerResultPartitionLocation(JobId jobId, boolean orderedResult, ResultSetId rsId,
            int partition, int nPartitions, NetworkAddress networkAddress) {
        Map<ResultSetId, Pair<Boolean, DatasetDirectoryRecord[]>> rsMap = jobResultLocationsMap.get(jobId);
        if (rsMap == null) {
            rsMap = new HashMap<ResultSetId, Pair<Boolean, DatasetDirectoryRecord[]>>();
            jobResultLocationsMap.put(jobId, rsMap);
        }

        DatasetDirectoryRecord[] records = resultSetsMap.get(rsId);
        if (records == null) {
            records = new DatasetDirectoryRecord[nPartitions];
            resultSetsMap.put(rsId, records);
        }

        if (records[partition] == null) {
            records[partition] = new DatasetDirectoryRecord();
        }
        records[partition].setNetworkAddress(networkAddress);
        notifyAll();
    }

    @Override
    public synchronized DatasetDirectoryRecord[] getResultPartitionLocations(JobId jobId, ResultSetId rsId,
            DatasetDirectoryRecord[] knownRecords) throws HyracksDataException {
        while (!newRecords(jobId, rsId, knownRecords)) {
            try {
                wait();
            } catch (InterruptedException e) {
                throw new HyracksDataException(e);
            }
        }
        return jobResultLocationsMap.get(jobId).get(rsId);
    }

    private boolean newRecords(JobId jobId, ResultSetId rsId, DatasetDirectoryRecord[] knownRecords) {
        Map<ResultSetId, DatasetDirectoryRecord[]> rsMap = jobResultLocationsMap.get(jobId);
        if (rsMap == null) {
            return false;
        }
        DatasetDirectoryRecord[] records = rsMap.get(rsId);
        return !Arrays.equals(records, knownRecords);
    }
}
@@ -752,16 +752,61 @@
 l.Map;%0A%0A
+import org.apache.commons.lang3.tuple.Pair;%0A%0A
 import e
@@ -1744,32 +1744,46 @@
 Map%3CResultSetId,
+ Pair%3CBoolean,
  DatasetDirector
@@ -1793,16 +1793,17 @@
 cord%5B%5D%3E%3E
+%3E
  jobResu
@@ -1917,32 +1917,46 @@
 Map%3CResultSetId,
+ Pair%3CBoolean,
  DatasetDirector
@@ -1958,32 +1958,33 @@
 ectoryRecord%5B%5D%3E%3E
+%3E
 ();%0A    %7D%0A%0A    @
@@ -2473,32 +2473,46 @@
       %7D%0A%0A       
+ Pair%3CBoolean,
  DatasetDirector
@@ -2520,26 +2520,17 @@
 Record%5B%5D
- records =
+%3E
  resultS
@@ -2527,24 +2527,32 @@
 %5D%3E resultSet
+Pair = r
 sMap.get(rsI
@@ -2569,21 +2569,27 @@
   if (re
-cords
+sultSetPair
  == null
@@ -2606,23 +2606,108 @@
       re
-cords =
+sultSetPair = Pair.%3CBoolean, DatasetDirectoryRecord%5B%5D%3E of(orderedResult,%0A                   
  new Dat
@@ -2738,16 +2738,17 @@
 titions%5D
+)
 ;%0A      
@@ -2754,24 +2754,16 @@
        r
-esultSet
 sMap.put
@@ -2771,21 +2771,27 @@
 rsId, re
-cords
+sultSetPair
 );%0A     
@@ -2788,32 +2788,101 @@
 ir);%0A        %7D%0A%0A
+        DatasetDirectoryRecord%5B%5D records = resultSetPair.getRight();%0A
         if (reco
@@ -3547,24 +3547,35 @@
 d).get(rsId)
+.getRight()
 ;%0A    %7D%0A%0A   
@@ -3671,16 +3671,56 @@
 Records)
+%0A            throws HyracksDataException
  %7B%0A     
@@ -3738,16 +3738,30 @@
 ltSetId,
+ Pair%3CBoolean,
  Dataset
@@ -3778,16 +3778,17 @@
 ecord%5B%5D%3E
+%3E
  rsMap =
@@ -3886,32 +3886,46 @@
        %7D%0A       
+ Pair%3CBoolean,
  DatasetDirector
@@ -3929,32 +3929,39 @@
 toryRecord%5B%5D
+%3E
  re
-cords
+sultSetPair
  = rsMap.get
@@ -3960,32 +3960,235 @@
 sMap.get(rsId);%0A
+        if (resultSetPair == null %7C%7C resultSetPair.getRight() == null) %7B%0A            throw new HyracksDataException(%22ResultSet locations uninitialized when it is expected to be initialized.%22);%0A        %7D%0A
         return !
@@ -4203,21 +4203,38 @@
 quals(re
-cords
+sultSetPair.getRight()
 , knownR

Use a slightly enhanced datastructure to store whether the results should be ordered or not.