package org.rakam.util.lock;

import com.google.common.base.Throwables;
import org.rakam.analysis.JDBCPoolDataSource;
import org.skife.jdbi.v2.DBI;
import org.skife.jdbi.v2.Handle;
import org.skife.jdbi.v2.util.BooleanMapper;

import java.sql.SQLException;
import java.util.Set;
import java.util.concurrent.ConcurrentSkipListSet;

public class MysqlLockService
        implements LockService {
    private final DBI dbi;
    private Handle currentHandle;
    private Set<String> locks;

    public MysqlLockService(JDBCPoolDataSource poolDataSource) {
        this.dbi = new DBI(() -> {
            return poolDataSource.getConnection(true);
        });
        this.currentHandle = dbi.open();
        locks = new ConcurrentSkipListSet<>();
    }

    @Override
    public Lock tryLock(String name) {
        if (!locks.add(name)) {
            return null;
        }
        try {
            return tryLock(name, 4);
        } catch (Exception e) {
            locks.remove(name);
            throw Throwables.propagate(e);
        }
    }

    public Lock tryLock(String name, int tryCount) {
        try {
            if (currentHandle.getConnection().isClosed()) {
                synchronized (this) {
                    currentHandle = dbi.open();
                }
            }

            Boolean first = currentHandle.createQuery("select get_lock(:name, 1)")
                    .bind("name", name)
                    .map(BooleanMapper.FIRST)
                    .first();

            if (!Boolean.TRUE.equals(first)) {
                locks.remove(name);
                return null;
            }

            return () -> {
                locks.remove(name);
                currentHandle.createQuery("select release_lock(:name)")
                        .bind("name", name)
                        .map(BooleanMapper.FIRST)
                        .first();
            };
        } catch (SQLException e) {
            try {
                if (currentHandle.getConnection().isClosed()) {
                    synchronized (this) {
                        currentHandle = dbi.open();
                    }
                }
            } catch (SQLException e1) {
                synchronized (this) {
                    currentHandle = dbi.open();
                }
            }

            if (tryCount == 0) {
                throw Throwables.propagate(e);
            }
            return tryLock(name, tryCount - 1);
        }
    }
}

@@ -224,38 +224,8 @@
 r;%0A%0A
-import java.sql.SQLException;%0A
 impo
@@ -1072,158 +1072,21 @@
     
-try %7B%0A            if (currentHandle.getConnection().isClosed()) %7B%0A                synchronized (this) %7B%0A                    currentHandle = dbi.open()
+Boolean first
 ;%0A  
@@ -1087,33 +1087,29 @@
 st;%0A        
-        %7D
+try %7B
 %0A           
@@ -1113,31 +1113,8 @@
     
-%7D%0A%0A            Boolean 
 firs
@@ -1291,414 +1291,8 @@
 t();
-%0A%0A            if (!Boolean.TRUE.equals(first)) %7B%0A                locks.remove(name);%0A                return null;%0A            %7D%0A%0A            return () -%3E %7B%0A                locks.remove(name);%0A                currentHandle.createQuery(%22select release_lock(:name)%22)%0A                        .bind(%22name%22, name)%0A                        .map(BooleanMapper.FIRST)%0A                        .first();%0A            %7D;
 %0A   
@@ -1297,35 +1297,32 @@
        %7D catch (
-SQL
 Exception e) %7B%0A 
@@ -1324,98 +1324,8 @@
 ) %7B%0A
-            try %7B%0A                if (currentHandle.getConnection().isClosed()) %7B%0A        
     
@@ -1366,32 +1366,24 @@
             
-   
-     
 currentHandl
@@ -1402,192 +1402,8 @@
 ();%0A
-                    %7D%0A                %7D%0A            %7D catch (SQLException e1) %7B%0A                synchronized (this) %7B%0A                    currentHandle = dbi.open();%0A                %7D%0A
     
@@ -1499,32 +1499,33 @@
 ;%0A            %7D%0A
+%0A
             retu
@@ -1561,21 +1561,383 @@
 );%0A        %7D
+%0A%0A        if (!Boolean.TRUE.equals(first)) %7B%0A            locks.remove(name);%0A            return null;%0A        %7D%0A%0A        return () -%3E %7B%0A            locks.remove(name);%0A            currentHandle.createQuery(%22select release_lock(:name)%22)%0A                    .bind(%22name%22, name)%0A                    .map(BooleanMapper.FIRST)%0A                    .first();%0A        %7D;
 %0A    %7D%0A%7D%0A

Fix bug in Mysql lock service