core/store/dist/src/main/java/org/onlab/onos/store/service/impl/DatabaseEntryExpirationTracker.java

/*
 * Copyright 2014 Open Networking Laboratory
 *
 * 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 org.onlab.onos.store.service.impl;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import net.jodah.expiringmap.ExpiringMap;
import net.jodah.expiringmap.ExpiringMap.ExpirationListener;
import net.jodah.expiringmap.ExpiringMap.ExpirationPolicy;
import net.kuujo.copycat.cluster.Member;
import net.kuujo.copycat.event.EventHandler;
import net.kuujo.copycat.event.LeaderElectEvent;

import org.onlab.onos.cluster.ClusterService;
import org.onlab.onos.store.cluster.messaging.ClusterCommunicationService;
import org.onlab.onos.store.cluster.messaging.ClusterMessage;
import org.onlab.onos.store.cluster.messaging.MessageSubject;
import org.onlab.onos.store.service.DatabaseService;
import org.onlab.onos.store.service.VersionedValue;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Plugs into the database update stream and track the TTL of entries added to
 * the database. For tables with pre-configured finite TTL, this class has
 * mechanisms for expiring (deleting) old, expired entries from the database.
 */
public class DatabaseEntryExpirationTracker implements
        DatabaseUpdateEventListener, EventHandler<LeaderElectEvent> {

    private final Logger log = LoggerFactory.getLogger(getClass());

    public static final MessageSubject DATABASE_UPDATES = new MessageSubject(
            "database-update-event");

    private DatabaseService databaseService;
    private ClusterService cluster;
    private ClusterCommunicationService clusterCommunicator;

    private final Member localMember;
    private final AtomicBoolean isLocalMemberLeader = new AtomicBoolean(false);

    private final Map<String, Map<DatabaseRow, VersionedValue>> tableEntryExpirationMap = new HashMap<>();

    private final ExpirationListener<DatabaseRow, VersionedValue> expirationObserver = new ExpirationObserver();

    DatabaseEntryExpirationTracker(Member localMember) {
        this.localMember = localMember;
    }

    @Override
    public void tableModified(TableModificationEvent event) {
        DatabaseRow row = new DatabaseRow(event.tableName(), event.key());
        Map<DatabaseRow, VersionedValue> map = tableEntryExpirationMap
                .get(event.tableName());

        switch (event.type()) {
        case ROW_DELETED:
            if (isLocalMemberLeader.get()) {
                try {
                    clusterCommunicator.broadcast(new ClusterMessage(cluster
                            .getLocalNode().id(), DATABASE_UPDATES,
                            DatabaseStateMachine.SERIALIZER.encode(event)));
                } catch (IOException e) {
                    log.error(
                            "Failed to broadcast a database table modification event.",
                            e);
                }
            }
            break;
        case ROW_ADDED:
        case ROW_UPDATED:
            map.put(row, null);
            break;
        default:
            break;
        }
    }

    @Override
    public void tableCreated(String tableName, int expirationTimeMillis) {
        // make this explicit instead of relying on a negative value
        // to indicate no expiration.
        if (expirationTimeMillis > 0) {
            tableEntryExpirationMap.put(tableName, ExpiringMap.builder()
                    .expiration(expirationTimeMillis, TimeUnit.SECONDS)
                    .expirationListener(expirationObserver)
                    // FIXME: make the expiration policy configurable.
                    .expirationPolicy(ExpirationPolicy.CREATED).build());
        }
    }

    @Override
    public void tableDeleted(String tableName) {
        tableEntryExpirationMap.remove(tableName);
    }

    private class ExpirationObserver implements
            ExpirationListener<DatabaseRow, VersionedValue> {
        @Override
        public void expired(DatabaseRow key, VersionedValue value) {
            try {
                if (isLocalMemberLeader.get()) {
                    if (!databaseService.removeIfVersionMatches(key.tableName,
                            key.key, value.version())) {
                        log.info("Entry in the database changed before right its TTL expiration.");
                    }
                } else {
                    // If this node is not the current leader, we should never
                    // let the expiring entries drop off
                    // Under stable conditions (i.e no leadership switch) the
                    // current leader will initiate
                    // a database remove and this instance will get notified
                    // of a tableModification event causing it to remove from
                    // the map.
                    Map<DatabaseRow, VersionedValue> map = tableEntryExpirationMap
                            .get(key.tableName);
                    if (map != null) {
                        map.put(key, value);
                    }
                }

            } catch (Exception e) {
                log.warn(
                        "Failed to delete entry from the database after ttl expiration. Will retry eviction",
                        e);
                tableEntryExpirationMap.get(key.tableName).put(
                        new DatabaseRow(key.tableName, key.key), value);
            }
        }
    }

    @Override
    public void handle(LeaderElectEvent event) {
        if (localMember.equals(event.leader())) {
            isLocalMemberLeader.set(true);
        }
    }

    /**
     * Wrapper class for a database row identifier.
     */
    private class DatabaseRow {

        String tableName;
        String key;

        public DatabaseRow(String tableName, String key) {
            this.tableName = tableName;
            this.key = key;
        }

        @Override
        public boolean equals(Object obj) {
            if (this == obj) {
                return true;
            }
            if (!(obj instanceof DatabaseRow)) {
                return false;
            }
            DatabaseRow that = (DatabaseRow) obj;

            return Objects.equals(this.tableName, that.tableName)
                    && Objects.equals(this.key, that.key);
        }

        @Override
        public int hashCode() {
            return Objects.hash(tableName, key);
        }
    }
}