<commit_before>/*
 * Copyright 2009 New Atlanta Communications, LLC
 *
 * 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 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.
 */
package com.newatlanta.appengine.datastore;

import static com.google.appengine.api.datastore.KeyFactory.keyToString;
import static com.google.appengine.api.labs.taskqueue.TaskOptions.Builder.url;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.logging.Logger;

import javax.servlet.ServletException;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import com.google.appengine.api.datastore.DatastoreService;
import com.google.appengine.api.datastore.DatastoreServiceFactory;
import com.google.appengine.api.datastore.DatastoreTimeoutException;
import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.EntityNotFoundException;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.datastore.KeyRange;
import com.google.appengine.api.datastore.PreparedQuery;
import com.google.appengine.api.datastore.Query;
import com.google.appengine.api.datastore.Transaction;
import com.google.appengine.api.labs.taskqueue.Queue;
import com.google.appengine.api.labs.taskqueue.QueueFactory;
import com.google.appengine.api.labs.taskqueue.TaskOptions;
import com.google.appengine.api.memcache.Expiration;
import com.google.appengine.api.memcache.MemcacheService;
import com.google.appengine.api.memcache.MemcacheServiceFactory;
import com.google.appengine.api.memcache.MemcacheService.SetPolicy;

/**
 * Implements a <code>DatastoreServer</code> that automatically caches entities
 * in memcache. Uses a write-behind cache when writing entities to the datastore.
 * 
 * @author <a href="mailto:<EMAIL>"><NAME></a>
 */
@SuppressWarnings("serial")
public class CachingDatastoreService extends HttpServlet implements DatastoreService {
    
    private static final String QUEUE_NAME = "write-behind-task";
    private static final String KEY_PARAM = "key";
    private static final String TASK_URL = "/" + QUEUE_NAME;
    
    private static DatastoreService datastore = DatastoreServiceFactory.getDatastoreService();
    private static MemcacheService memcache = MemcacheServiceFactory.getMemcacheService();
    private static Queue queue;
    
    private static final Logger log = Logger.getLogger( CachingDatastoreService.class.getName() );
    
    static {
        try {
            queue = QueueFactory.getQueue( QUEUE_NAME );
            queue.add( url( TASK_URL ) ); // verify queue configured
        } catch ( Exception e ) {
            log.info( e.getMessage() );
            queue = QueueFactory.getDefaultQueue(); // TODO: test default queue
        }
    }
    
    private Expiration expiration;
    
    public CachingDatastoreService() {
        this( null );
    }
    
    public CachingDatastoreService( Expiration expiration ) {
        this.expiration = expiration;
    }
    
    /**
     * Gets an entity. Returns the entity from memcache, if it exists; otherwise,
     * gets the entity from the datastore, puts it into memcache, then returns it.
     */
    public Entity get( Key key ) throws EntityNotFoundException {
        return get( null, key );
    }
    
    public Entity get( Transaction txn, Key key ) throws EntityNotFoundException {
        String keyString = keyToString( key );
        Entity entity = (Entity)memcache.get( keyString );
        if ( entity == null ) {
            try {
                entity = datastore.get( txn, key );
            } catch ( DatastoreTimeoutException e ) {
                entity = datastore.get( txn, key );
            }
            memcache.put( keyString, entity, expiration, SetPolicy.SET_ALWAYS );
        }
        return entity;
    }
    
    /**
     * WARNING! Don't rely on the ordering of entities returned by this method.
     */
    public Map<Key, Entity> get( Iterable<Key> keys ) {
        return get( null, keys );
    }
    
    
    @SuppressWarnings("unchecked")
    public Map<Key, Entity> get( Transaction txn, Iterable<Key> keys ) {
        // TODO: this method has not been tested
        Collection<Object> keyStrings = getKeyStrings( keys );
        Map<Key, Entity> datastoreMap = getDatastoreMap( (Map)memcache.getAll( keyStrings ) );
        if ( datastoreMap.isEmpty() ) {
            return getAndCache( txn, keys );
        }
        if ( datastoreMap.size() < keyStrings.size() ) {
            List<Key> notFound = new ArrayList<Key>();
            for ( Key key : keys ) {
                if ( !datastoreMap.containsKey( key ) ) {
                    notFound.add( key );
                }
            }
            if ( notFound.size() > 0 ) {
                try {
                    datastoreMap.putAll( getAndCache( txn, notFound ) );
                } catch ( Exception e ) {
                    log.warning( e.getMessage() );
                }
            }
        }
        return datastoreMap;
    }
    
    private static Collection<Object> getKeyStrings( Iterable<Key> keys ) {
        Collection<Object> keyStrings = new ArrayList<Object>();
        for ( Key key : keys ) {
            keyStrings.add( keyToString( key ) );
        }
        return keyStrings;
    }
    
    private static Map<Key, Entity> getDatastoreMap( Map<String, Entity> memcacheMap ) {
        Map<Key, Entity> datastoreMap = new HashMap<Key, Entity>();
        for ( Entity entity : memcacheMap.values() ) {
            datastoreMap.put( entity.getKey(), entity );
        }
        return datastoreMap;
    }

    @SuppressWarnings("unchecked")
    private Map<Key, Entity> getAndCache( Transaction txn, Iterable<Key> keys ) {
        Map<Key, Entity> entityMap;
        try {
            entityMap = datastore.get( txn, keys );
        } catch ( DatastoreTimeoutException e ) {
            entityMap = datastore.get( txn, keys );
        }
        memcache.putAll( (Map)getMemcacheMap( entityMap ), expiration, SetPolicy.SET_ALWAYS );
        return entityMap;
    }
    
    private static Map<String, Entity> getMemcacheMap( Map<Key, Entity> datastoreMap ) {
        Map<String, Entity> memcacheMap = new HashMap<String, Entity>();
        for ( Entity entity : datastoreMap.values() ) {
            memcacheMap.put( keyToString( entity.getKey() ), entity );
        }
        return memcacheMap;
    }
    
    public Key put( Entity entity ) {
        Key key = entity.getKey();
        if ( !key.isComplete() ) {
            // TODO: this path has not been tested
            KeyRange keyRange = datastore.allocateIds( key.getParent(), key.getKind(), 1 );
            key = keyRange.getStart();
        }
        String keyString = keyToString( key );
        memcache.put( keyString, entity, expiration, SetPolicy.SET_ALWAYS );
        try {
            queue.add( url( TASK_URL ).param( KEY_PARAM, keyString ) );
            log.info( key.toString() );
        } catch ( Exception e ) {
            log.warning( e.getMessage() );
            try {
                key = datastore.put( entity );
            } catch ( DatastoreTimeoutException t ) {
                key = datastore.put( entity );
            }
        }
        return key;
    }
    
    @SuppressWarnings("unchecked")
    public List<Key> put( Iterable<Entity> entities ) {
        List<Key> keyList = new ArrayList<Key>();
        Map<String, Entity> memcacheMap = new HashMap<String, Entity>();
        TaskOptions taskOptions = url( TASK_URL );
        int i = 0;
        for ( Entity entity : entities ) {
            Key key = entity.getKey();
            if ( !key.isComplete() ) {
                // TODO: needs to be implemented
            }
            keyList.add( key );
            String keyString = keyToString( key );
            memcacheMap.put( keyString, entity );
            taskOptions.param( KEY_PARAM + i++, keyString );
            log.info( key.toString() );
        }
        memcache.putAll( (Map)memcacheMap, expiration, SetPolicy.SET_ALWAYS );
        try {
            queue.add( taskOptions );
            return keyList;
        } catch ( Exception e ) {
            log.warning( e.getMessage() );
            try {
                return datastore.put( entities );
            } catch ( DatastoreTimeoutException t ) {
                return datastore.put( entities );
            }
        }
    }
    
    /**
     * Don't use write-behind cache with transactions.
     */
    public Key put( Transaction txn, Entity entity ) {
        return datastore.put( txn, entity );
    }

    public List<Key> put( Transaction txn, Iterable<Entity> entities ) {
        return datastore.put( txn, entities );
    }
    
    public void delete( Key ... keys ) {
        delete( null, Arrays.asList( keys ) );
    }
    
    public void delete( Iterable<Key> keys ) {
        delete( null, keys );
    }
    
    public void delete( Transaction txn, Key ... keys ) {
        delete( txn, Arrays.asList( keys ) ); 
    }

    @SuppressWarnings("unchecked")
    public void delete( Transaction txn, Iterable<Key> keys ) {
        try {
            datastore.delete( txn, keys );
        } catch ( DatastoreTimeoutException e ) {
            datastore.delete( txn, keys );
        }
        memcache.deleteAll( (Collection)keys );
    }

    public KeyRange allocateIds( String kind, long num ) {
        return datastore.allocateIds( kind, num );
    }

    public KeyRange allocateIds( Key parent, String kind, long num ) {
        return datastore.allocateIds( parent, kind, num );
    }

    public Transaction beginTransaction() {
        return datastore.beginTransaction();
    }

    public Collection<Transaction> getActiveTransactions() {
        return datastore.getActiveTransactions();
    }

    public Transaction getCurrentTransaction() {
        return datastore.getCurrentTransaction();
    }

    public Transaction getCurrentTransaction( Transaction returnedIfNoTxn ) {
        return datastore.getCurrentTransaction( returnedIfNoTxn );
    }

    // TODO: is there a way to cache query results?
    public PreparedQuery prepare( Query query ) {
        return datastore.prepare( query );
    }

    public PreparedQuery prepare( Transaction txn, Query query ) {
        return datastore.prepare( txn, query );
    }
    
    /***************************************************************************
     *                                                                         *
     *               W R I T E - B E H I N D   T A S K                         *
     *                                                                         *
     ***************************************************************************/
    
    @Override
    public void doGet( HttpServletRequest req, HttpServletResponse res )
            throws ServletException, IOException {
        doWriteBehindTask( req, res );
    }

    @Override
    public void doPost( HttpServletRequest req, HttpServletResponse res )
            throws ServletException, IOException {
        doWriteBehindTask( req, res );
    }

    @SuppressWarnings("unchecked")
    private static void doWriteBehindTask( HttpServletRequest req, HttpServletResponse res )
            throws IOException {
        List<String> keys = new ArrayList<String>();
        Enumeration<String> paramNames = req.getParameterNames();
        while ( paramNames.hasMoreElements() ) {
            // parameters can have only a single value (see issue #2090)
            keys.add( req.getParameter( paramNames.nextElement() ) );
        }
        Map<String, Entity> entityMap = (Map)memcache.getAll( (List)keys );
        if ( !entityMap.isEmpty() ) {
            // TODO: using transactions should be optional (see issue #2090)
            Transaction txn = datastore.beginTransaction();
            try {
                datastore.put( txn, entityMap.values() );
                txn.commit();
            } catch ( DatastoreTimeoutException e ) { // retry task
                res.sendError( HttpServletResponse.SC_REQUEST_TIMEOUT );
            } catch ( Exception e ) {
                log.warning( e.getMessage() ); // don't retry
            } finally {
                if ( txn.isActive() ) {
                    txn.rollback();
                }
            }
        }
    }
}
<commit_msg>corrected error handling<commit_after>/*
 * Copyright 2009 New Atlanta Communications, LLC
 *
 * 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 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.
 */
package com.newatlanta.appengine.datastore;

import static com.google.appengine.api.datastore.KeyFactory.keyToString;
import static com.google.appengine.api.labs.taskqueue.TaskOptions.Builder.url;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.logging.Logger;

import javax.servlet.ServletException;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import com.google.appengine.api.datastore.DatastoreService;
import com.google.appengine.api.datastore.DatastoreServiceFactory;
import com.google.appengine.api.datastore.DatastoreTimeoutException;
import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.EntityNotFoundException;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.datastore.KeyRange;
import com.google.appengine.api.datastore.PreparedQuery;
import com.google.appengine.api.datastore.Query;
import com.google.appengine.api.datastore.Transaction;
import com.google.appengine.api.labs.taskqueue.Queue;
import com.google.appengine.api.labs.taskqueue.QueueFactory;
import com.google.appengine.api.labs.taskqueue.TaskOptions;
import com.google.appengine.api.memcache.Expiration;
import com.google.appengine.api.memcache.MemcacheService;
import com.google.appengine.api.memcache.MemcacheServiceFactory;
import com.google.appengine.api.memcache.MemcacheService.SetPolicy;

/**
 * Implements a <code>DatastoreServer</code> that automatically caches entities
 * in memcache. Uses a write-behind cache when writing entities to the datastore.
 * 
 * @author <a href="mailto:<EMAIL>"><NAME></a>
 */
@SuppressWarnings("serial")
public class CachingDatastoreService extends HttpServlet implements DatastoreService {
    
    private static final String QUEUE_NAME = "write-behind-task";
    private static final String KEY_PARAM = "key";
    private static final String TASK_URL = "/" + QUEUE_NAME;
    
    private static DatastoreService datastore = DatastoreServiceFactory.getDatastoreService();
    private static MemcacheService memcache = MemcacheServiceFactory.getMemcacheService();
    private static Queue queue;
    
    private static final Logger log = Logger.getLogger( CachingDatastoreService.class.getName() );
    
    static {
        try {
            queue = QueueFactory.getQueue( QUEUE_NAME );
            queue.add( url( TASK_URL ) ); // verify queue configured
        } catch ( Exception e ) {
            log.info( e.getMessage() );
            queue = QueueFactory.getDefaultQueue(); // TODO: test default queue
        }
    }
    
    private Expiration expiration;
    
    public CachingDatastoreService() {
        this( null );
    }
    
    public CachingDatastoreService( Expiration expiration ) {
        this.expiration = expiration;
    }
    
    /**
     * Gets an entity. Returns the entity from memcache, if it exists; otherwise,
     * gets the entity from the datastore, puts it into memcache, then returns it.
     */
    public Entity get( Key key ) throws EntityNotFoundException {
        return get( null, key );
    }
    
    public Entity get( Transaction txn, Key key ) throws EntityNotFoundException {
        String keyString = keyToString( key );
        Entity entity = (Entity)memcache.get( keyString );
        if ( entity == null ) {
            try {
                entity = datastore.get( txn, key );
            } catch ( DatastoreTimeoutException e ) {
                entity = datastore.get( txn, key );
            }
            memcache.put( keyString, entity, expiration, SetPolicy.SET_ALWAYS );
        }
        return entity;
    }
    
    /**
     * WARNING! Don't rely on the ordering of entities returned by this method.
     */
    public Map<Key, Entity> get( Iterable<Key> keys ) {
        return get( null, keys );
    }
    
    
    @SuppressWarnings("unchecked")
    public Map<Key, Entity> get( Transaction txn, Iterable<Key> keys ) {
        // TODO: this method has not been tested
        Collection<Object> keyStrings = getKeyStrings( keys );
        Map<Key, Entity> datastoreMap = getDatastoreMap( (Map)memcache.getAll( keyStrings ) );
        if ( datastoreMap.isEmpty() ) {
            return getAndCache( txn, keys );
        }
        if ( datastoreMap.size() < keyStrings.size() ) {
            List<Key> notFound = new ArrayList<Key>();
            for ( Key key : keys ) {
                if ( !datastoreMap.containsKey( key ) ) {
                    notFound.add( key );
                }
            }
            datastoreMap.putAll( getAndCache( txn, notFound ) );
        }
        return datastoreMap;
    }
    
    private static Collection<Object> getKeyStrings( Iterable<Key> keys ) {
        Collection<Object> keyStrings = new ArrayList<Object>();
        for ( Key key : keys ) {
            keyStrings.add( keyToString( key ) );
        }
        return keyStrings;
    }
    
    private static Map<Key, Entity> getDatastoreMap( Map<String, Entity> memcacheMap ) {
        Map<Key, Entity> datastoreMap = new HashMap<Key, Entity>();
        for ( Entity entity : memcacheMap.values() ) {
            datastoreMap.put( entity.getKey(), entity );
        }
        return datastoreMap;
    }

    @SuppressWarnings("unchecked")
    private Map<Key, Entity> getAndCache( Transaction txn, Iterable<Key> keys ) {
        Map<Key, Entity> entityMap;
        try {
            entityMap = datastore.get( txn, keys );
        } catch ( DatastoreTimeoutException e ) {
            entityMap = datastore.get( txn, keys );
        }
        memcache.putAll( (Map)getMemcacheMap( entityMap ), expiration, SetPolicy.SET_ALWAYS );
        return entityMap;
    }
    
    private static Map<String, Entity> getMemcacheMap( Map<Key, Entity> datastoreMap ) {
        Map<String, Entity> memcacheMap = new HashMap<String, Entity>();
        for ( Entity entity : datastoreMap.values() ) {
            memcacheMap.put( keyToString( entity.getKey() ), entity );
        }
        return memcacheMap;
    }
    
    public Key put( Entity entity ) {
        Key key = entity.getKey();
        if ( !key.isComplete() ) {
            // TODO: this path has not been tested
            KeyRange keyRange = datastore.allocateIds( key.getParent(), key.getKind(), 1 );
            key = keyRange.getStart();
        }
        String keyString = keyToString( key );
        memcache.put( keyString, entity, expiration, SetPolicy.SET_ALWAYS );
        try {
            queue.add( url( TASK_URL ).param( KEY_PARAM, keyString ) );
            log.info( key.toString() );
        } catch ( Exception e ) {
            log.warning( e.getMessage() );
            try {
                key = datastore.put( entity );
            } catch ( DatastoreTimeoutException t ) {
                key = datastore.put( entity );
            }
        }
        return key;
    }
    
    @SuppressWarnings("unchecked")
    public List<Key> put( Iterable<Entity> entities ) {
        List<Key> keyList = new ArrayList<Key>();
        Map<String, Entity> memcacheMap = new HashMap<String, Entity>();
        TaskOptions taskOptions = url( TASK_URL );
        int i = 0;
        for ( Entity entity : entities ) {
            Key key = entity.getKey();
            if ( !key.isComplete() ) {
                // TODO: needs to be implemented
            }
            keyList.add( key );
            String keyString = keyToString( key );
            memcacheMap.put( keyString, entity );
            taskOptions.param( KEY_PARAM + i++, keyString );
            log.info( key.toString() );
        }
        memcache.putAll( (Map)memcacheMap, expiration, SetPolicy.SET_ALWAYS );
        try {
            queue.add( taskOptions );
            return keyList;
        } catch ( Exception e ) {
            log.warning( e.getMessage() );
            try {
                return datastore.put( entities );
            } catch ( DatastoreTimeoutException t ) {
                return datastore.put( entities );
            }
        }
    }
    
    /**
     * Don't use write-behind cache with transactions.
     */
    public Key put( Transaction txn, Entity entity ) {
        return datastore.put( txn, entity );
    }

    public List<Key> put( Transaction txn, Iterable<Entity> entities ) {
        return datastore.put( txn, entities );
    }
    
    public void delete( Key ... keys ) {
        delete( null, Arrays.asList( keys ) );
    }
    
    public void delete( Iterable<Key> keys ) {
        delete( null, keys );
    }
    
    public void delete( Transaction txn, Key ... keys ) {
        delete( txn, Arrays.asList( keys ) ); 
    }

    @SuppressWarnings("unchecked")
    public void delete( Transaction txn, Iterable<Key> keys ) {
        try {
            datastore.delete( txn, keys );
        } catch ( DatastoreTimeoutException e ) {
            datastore.delete( txn, keys );
        }
        memcache.deleteAll( (Collection)keys );
    }

    public KeyRange allocateIds( String kind, long num ) {
        return datastore.allocateIds( kind, num );
    }

    public KeyRange allocateIds( Key parent, String kind, long num ) {
        return datastore.allocateIds( parent, kind, num );
    }

    public Transaction beginTransaction() {
        return datastore.beginTransaction();
    }

    public Collection<Transaction> getActiveTransactions() {
        return datastore.getActiveTransactions();
    }

    public Transaction getCurrentTransaction() {
        return datastore.getCurrentTransaction();
    }

    public Transaction getCurrentTransaction( Transaction returnedIfNoTxn ) {
        return datastore.getCurrentTransaction( returnedIfNoTxn );
    }

    // TODO: is there a way to cache query results?
    public PreparedQuery prepare( Query query ) {
        return datastore.prepare( query );
    }

    public PreparedQuery prepare( Transaction txn, Query query ) {
        return datastore.prepare( txn, query );
    }
    
    /***************************************************************************
     *                                                                         *
     *               W R I T E - B E H I N D   T A S K                         *
     *                                                                         *
     ***************************************************************************/
    
    @Override
    public void doGet( HttpServletRequest req, HttpServletResponse res )
            throws ServletException, IOException {
        doWriteBehindTask( req, res );
    }

    @Override
    public void doPost( HttpServletRequest req, HttpServletResponse res )
            throws ServletException, IOException {
        doWriteBehindTask( req, res );
    }

    @SuppressWarnings("unchecked")
    private static void doWriteBehindTask( HttpServletRequest req, HttpServletResponse res )
            throws IOException {
        List<String> keys = new ArrayList<String>();
        Enumeration<String> paramNames = req.getParameterNames();
        while ( paramNames.hasMoreElements() ) {
            // parameters can have only a single value (see issue #2090)
            keys.add( req.getParameter( paramNames.nextElement() ) );
        }
        Map<String, Entity> entityMap = (Map)memcache.getAll( (List)keys );
        if ( !entityMap.isEmpty() ) {
            // TODO: using transactions should be optional (see issue #2090)
            Transaction txn = datastore.beginTransaction();
            try {
                datastore.put( txn, entityMap.values() );
                txn.commit();
            } catch ( DatastoreTimeoutException e ) { // retry task
                res.sendError( HttpServletResponse.SC_REQUEST_TIMEOUT );
            } catch ( Exception e ) {
                log.warning( e.getMessage() ); // don't retry
            } finally {
                if ( txn.isActive() ) {
                    txn.rollback();
                }
            }
        }
    }
}

{'fraction_non_alphanumeric': 0.10596750480642364, 'fraction_numerical': 0.0013948052927206243, 'mean_word_length': 2.075710144927536, 'pattern_counts': {'":': 0, '<': 93, '<?xml version=': 0, '>': 92, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1116267', 'n_tokens_mistral': 7192, 'n_tokens_neox': 7011, 'n_words': 1833}