package org.apache.helix.controller.restlet;
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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 at
*
* 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.
*/
import java.io.StringWriter;
import java.util.Map;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.helix.ZNRecord;
import org.apache.log4j.Logger;
import org.codehaus.jackson.map.ObjectMapper;
import org.restlet.Client;
import org.restlet.data.MediaType;
import org.restlet.data.Method;
import org.restlet.data.Protocol;
import org.restlet.data.Reference;
import org.restlet.data.Request;
import org.restlet.data.Response;
import org.restlet.data.Status;
public class ZkPropertyTransferClient
{
private static Logger LOG = Logger.getLogger(ZkPropertyTransferClient.class);
public static final int DEFAULT_MAX_CONCURRENTTASKS = 2;
public static int SEND_PERIOD = 10 * 1000;
public static final String USE_PROPERTYTRANSFER = "UsePropertyTransfer";
int _maxConcurrentTasks;
ExecutorService _executorService;
Client[] _clients;
AtomicInteger _requestCount = new AtomicInteger(0);
// ZNRecord update buffer: key is the zkPath, value is the ZNRecordUpdate
AtomicReference<ConcurrentHashMap<String, ZNRecordUpdate>> _dataBufferRef
= new AtomicReference<ConcurrentHashMap<String, ZNRecordUpdate>>();
Timer _timer;
volatile String _webServiceUrl = "";
public ZkPropertyTransferClient(int maxConcurrentTasks)
{
_maxConcurrentTasks = maxConcurrentTasks;
_executorService = Executors.newFixedThreadPool(_maxConcurrentTasks);
_clients = new Client[_maxConcurrentTasks];
for(int i = 0; i< _clients.length; i++)
{
_clients[i] = new Client(Protocol.HTTP);
}
_timer = new Timer(true);
_timer.schedule(new SendZNRecordTimerTask(), SEND_PERIOD, SEND_PERIOD);
_dataBufferRef.getAndSet(new ConcurrentHashMap<String, ZNRecordUpdate>());
}
class SendZNRecordTimerTask extends TimerTask
{
@Override
public void run()
{
sendUpdateBatch();
}
}
public void enqueueZNRecordUpdate(ZNRecordUpdate update, String webserviceUrl)
{
try
{
LOG.info("Enqueue update to " + update.getPath() + " opcode: " + update.getOpcode() + " to " + webserviceUrl);
_webServiceUrl = webserviceUrl;
update.getRecord().setSimpleField(USE_PROPERTYTRANSFER, "true");
synchronized(_dataBufferRef)
{
if(_dataBufferRef.get().containsKey(update._path))
{
ZNRecord oldVal = _dataBufferRef.get().get(update.getPath()).getRecord();
oldVal = update.getZNRecordUpdater().update(oldVal);
_dataBufferRef.get().get(update.getPath())._record = oldVal;
}
else
{
_dataBufferRef.get().put(update.getPath(), update);
}
}
}
catch(Exception e)
{
LOG.error("", e);
}
}
void sendUpdateBatch()
{
LOG.debug("Actual sending update with " + _dataBufferRef.get().size() + " updates to " + _webServiceUrl);
Map<String, ZNRecordUpdate> updateCache = null;
synchronized(_dataBufferRef)
{
updateCache = _dataBufferRef.getAndSet(new ConcurrentHashMap<String, ZNRecordUpdate>());
}
if(updateCache != null && updateCache.size() > 0)
{
ZNRecordUpdateUploadTask task = new ZNRecordUpdateUploadTask(updateCache, _webServiceUrl, _clients[_requestCount.intValue() % _maxConcurrentTasks]);
_requestCount.incrementAndGet();
_executorService.submit(task);
LOG.trace("Queue size :" + ((ThreadPoolExecutor)_executorService).getQueue().size());
}
}
public void shutdown()
{
LOG.info("Shutting down ZkPropertyTransferClient");
_executorService.shutdown();
_timer.cancel();
for(Client client: _clients)
{
try
{
client.stop();
}
catch (Exception e)
{
LOG.error("", e);
}
}
}
class ZNRecordUpdateUploadTask implements Callable<Void>
{
Map<String, ZNRecordUpdate> _updateMap;
String _webServiceUrl;
Client _client;
ZNRecordUpdateUploadTask(Map<String, ZNRecordUpdate> update, String webserviceUrl, Client client)
{
_updateMap = update;
_webServiceUrl = webserviceUrl;
_client = client;
}
@Override
public Void call() throws Exception
{
LOG.debug("Actual sending update with " + _updateMap.size() + " updates to " + _webServiceUrl);
long time = System.currentTimeMillis();
Reference resourceRef = new Reference(_webServiceUrl);
Request request = new Request(Method.PUT, resourceRef);
ObjectMapper mapper = new ObjectMapper();
StringWriter sw = new StringWriter();
try
{
mapper.writeValue(sw, _updateMap);
}
catch (Exception e)
{
LOG.error("",e);
}
request.setEntity(
ZNRecordUpdateResource.UPDATEKEY + "=" + sw, MediaType.APPLICATION_ALL);
// This is a sync call. See com.noelios.restlet.http.StreamClientCall.sendRequest()
Response response = _client.handle(request);
if(response.getStatus().getCode() != Status.SUCCESS_OK.getCode())
{
LOG.error("Status : " + response.getStatus());
}
LOG.info("Using time : " + (System.currentTimeMillis() - time));
return null;
}
}
}