6918573: sun.security.pkcs11.P11RSACipher.finalize() is a scalability blocker

Summary: Removed the finalize() methods and use PhantomReference in Session to do auto clean up.
Reviewed-by: wetmore

src/share/classes/sun/security/pkcs11/P11Cipher.java

 /*
- * Copyright 2003-2008 Sun Microsystems, Inc.  All Rights Reserved.
+ * Copyright 2003-2010 Sun Microsystems, Inc.  All Rights Reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -192,7 +192,6 @@ final class P11Cipher extends CipherSpi {
             // should not happen
             throw new ProviderException(nspe);
         }
-        session = token.getOpSession();
     }
 
     protected void engineSetMode(String mode) throws NoSuchAlgorithmException {
@@ -847,18 +846,6 @@ final class P11Cipher extends CipherSpi {
         return n;
     }
 
-    @Override
-    protected void finalize() throws Throwable {
-        try {
-            if ((session != null) && token.isValid()) {
-                cancelOperation();
-                session = token.releaseSession(session);
-            }
-        } finally {
-            super.finalize();
-        }
-    }
-
     private final void bufferInputBytes(byte[] in, int inOfs, int len) {
         System.arraycopy(in, inOfs, padBuffer, padBufferLen, len);
         padBufferLen += len;

src/share/classes/sun/security/pkcs11/P11Digest.java

 /*
- * Copyright 2003-2005 Sun Microsystems, Inc.  All Rights Reserved.
+ * Copyright 2003-2010 Sun Microsystems, Inc.  All Rights Reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -308,16 +308,4 @@ final class P11Digest extends MessageDigestSpi {
             throw new ProviderException("update() failed", e);
         }
     }
-
-    protected void finalize() throws Throwable {
-        try {
-            if ((session != null) && token.isValid()) {
-                cancelOperation();
-                session = token.releaseSession(session);
-            }
-        } finally {
-            super.finalize();
-        }
-    }
-
 }

src/share/classes/sun/security/pkcs11/P11Key.java

 /*
- * Copyright 2003-2009 Sun Microsystems, Inc.  All Rights Reserved.
+ * Copyright 2003-2010 Sun Microsystems, Inc.  All Rights Reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -85,7 +85,7 @@ abstract class P11Key implements Key {
     // flags indicating whether the key is a token object, sensitive, extractable
     final boolean tokenObject, sensitive, extractable;
 
-    // weak reference notification clean up for session keys
+    // phantom reference notification clean up for session keys
     private final SessionKeyRef sessionKeyRef;
 
     P11Key(String type, Session session, long keyID, String algorithm,
@@ -1051,7 +1051,12 @@ abstract class P11Key implements Key {
     }
 }
 
-final class SessionKeyRef extends WeakReference<P11Key>
+/*
+ * NOTE: Must use PhantomReference here and not WeakReference
+ * otherwise the key maybe cleared before other objects which
+ * still use these keys during finalization such as SSLSocket.
+ */
+final class SessionKeyRef extends PhantomReference<P11Key>
     implements Comparable<SessionKeyRef> {
     private static ReferenceQueue<P11Key> refQueue =
         new ReferenceQueue<P11Key>();
@@ -1062,14 +1067,11 @@ final class SessionKeyRef extends WeakReference<P11Key>
         return refQueue;
     }
 
-    static final private int MAX_ITERATIONS = 2;
-
     private static void drainRefQueueBounded() {
-        int iterations = 0;
-        while (iterations < MAX_ITERATIONS) {
+        while (true) {
             SessionKeyRef next = (SessionKeyRef) refQueue.poll();
-            if (next != null) next.dispose();
-            ++iterations;
+            if (next == null) break;
+            next.dispose();
         }
     }
 
@@ -1087,7 +1089,7 @@ final class SessionKeyRef extends WeakReference<P11Key>
         drainRefQueueBounded();
     }
 
-    void dispose() {
+    private void dispose() {
         refList.remove(this);
         if (session.token.isValid()) {
             Session newSession = null;
@@ -1097,6 +1099,7 @@ final class SessionKeyRef extends WeakReference<P11Key>
             } catch (PKCS11Exception e) {
                 // ignore
             } finally {
+                this.clear();
                 session.token.releaseSession(newSession);
                 session.removeObject();
             }

src/share/classes/sun/security/pkcs11/P11Mac.java

 /*
- * Copyright 2003-2007 Sun Microsystems, Inc.  All Rights Reserved.
+ * Copyright 2003-2010 Sun Microsystems, Inc.  All Rights Reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -263,16 +263,4 @@ final class P11Mac extends MacSpi {
             throw new ProviderException("update() failed", e);
         }
     }
-
-    protected void finalize() throws Throwable {
-        try {
-            if ((session != null) && token.isValid()) {
-                cancelOperation();
-                session = token.releaseSession(session);
-            }
-        } finally {
-            super.finalize();
-        }
-    }
-
 }

src/share/classes/sun/security/pkcs11/P11RSACipher.java

 /*
- * Copyright 2003-2009 Sun Microsystems, Inc.  All Rights Reserved.
+ * Copyright 2003-2010 Sun Microsystems, Inc.  All Rights Reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -485,18 +485,6 @@ final class P11RSACipher extends CipherSpi {
         int n = P11KeyFactory.convertKey(token, key, algorithm).keyLength();
         return n;
     }
-
-    protected void finalize() throws Throwable {
-        try {
-            if ((session != null) && token.isValid()) {
-                cancelOperation();
-                session = token.releaseSession(session);
-            }
-        } finally {
-            super.finalize();
-        }
-    }
-
 }
 
 final class ConstructKeys {

src/share/classes/sun/security/pkcs11/P11Signature.java

@@ -226,7 +226,6 @@ final class P11Signature extends SignatureSpi {
         this.buffer = buffer;
         this.digestOID = digestOID;
         this.md = md;
-        session = token.getOpSession();
     }
 
     private void ensureInitialized() {
@@ -732,16 +731,4 @@ final class P11Signature extends SignatureSpi {
             throws InvalidParameterException {
         throw new UnsupportedOperationException("getParameter() not supported");
     }
-
-    protected void finalize() throws Throwable {
-        try {
-            if ((session != null) && token.isValid()) {
-                cancelOperation();
-                session = token.releaseSession(session);
-            }
-        } finally {
-            super.finalize();
-        }
-    }
-
 }

src/share/classes/sun/security/pkcs11/Session.java

 /*
- * Copyright 2003-2005 Sun Microsystems, Inc.  All Rights Reserved.
+ * Copyright 2003-2010 Sun Microsystems, Inc.  All Rights Reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -25,6 +25,7 @@
 
 package sun.security.pkcs11;
 
+import java.lang.ref.*;
 import java.util.*;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -59,11 +60,14 @@ final class Session implements Comparable<Session> {
     // this could lead to idle sessions being closed early, but that is harmless
     private long lastAccess;
 
+    private final SessionRef sessionRef;
+
     Session(Token token, long id) {
         this.token = token;
         this.id = id;
         createdObjects = new AtomicInteger();
         id();
+        sessionRef = new SessionRef(this, id, token);
     }
 
     public int compareTo(Session other) {
@@ -108,4 +112,76 @@ final class Session implements Comparable<Session> {
         return createdObjects.get() != 0;
     }
 
+    void close() {
+        if (hasObjects()) {
+            throw new ProviderException(
+                "Internal error: close session with active objects");
+        }
+        sessionRef.dispose();
+    }
+}
+
+/*
+ * NOTE: Use PhantomReference here and not WeakReference
+ * otherwise the sessions maybe closed before other objects
+ * which are still being finalized.
+ */
+final class SessionRef extends PhantomReference<Session>
+        implements Comparable<SessionRef> {
+
+    private static ReferenceQueue<Session> refQueue =
+        new ReferenceQueue<Session>();
+
+    private static Set<SessionRef> refList =
+        Collections.synchronizedSortedSet(new TreeSet<SessionRef>());
+
+    static ReferenceQueue<Session> referenceQueue() {
+        return refQueue;
+    }
+
+    static int totalCount() {
+        return refList.size();
+    }
+
+    private static void drainRefQueueBounded() {
+        while (true) {
+            SessionRef next = (SessionRef) refQueue.poll();
+            if (next == null) break;
+            next.dispose();
+        }
+    }
+
+    // handle to the native session
+    private long id;
+    private Token token;
+
+    SessionRef(Session session, long id, Token token) {
+        super(session, refQueue);
+        this.id = id;
+        this.token = token;
+        refList.add(this);
+        // TBD: run at some interval and not every time?
+        drainRefQueueBounded();
+    }
+
+    void dispose() {
+        refList.remove(this);
+        try {
+            token.p11.C_CloseSession(id);
+        } catch (PKCS11Exception e1) {
+            // ignore
+        } catch (ProviderException e2) {
+            // ignore
+        } finally {
+            this.clear();
+        }
+    }
+
+    public int compareTo(SessionRef other) {
+        if (this.id == other.id) {
+            return 0;
+        } else {
+            return (this.id < other.id) ? -1 : 1;
+        }
+    }
 }

src/share/classes/sun/security/pkcs11/SessionManager.java

 /*
- * Copyright 2003-2006 Sun Microsystems, Inc.  All Rights Reserved.
+ * Copyright 2003-2010 Sun Microsystems, Inc.  All Rights Reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -51,10 +51,12 @@ import static sun.security.pkcs11.wrapper.PKCS11Constants.*;
  * number of such sessions low. Note that we occasionally want to explicitly
  * close a session, see P11Signature.
  *
- * NOTE that all sessions obtained from this class MUST be returned using
- * either releaseSession() or closeSession() using a finally block or a
- * finalizer where appropriate. Otherwise, they will be "lost", i.e. there
- * will be a resource leak eventually leading to exhaustion.
+ * NOTE that sessions obtained from this class SHOULD be returned using
+ * either releaseSession() or closeSession() using a finally block when
+ * not needed anymore. Otherwise, they will be left for cleanup via the
+ * PhantomReference mechanism when GC kicks in, but it's best not to rely
+ * on that since GC may not run timely enough since the native PKCS11 library
+ * is also consuming memory.
  *
  * Note that sessions are automatically closed when they are not used for a
  * period of time, see Session.
@@ -74,9 +76,6 @@ final class SessionManager {
     // maximum number of sessions to open with this token
     private final int maxSessions;
 
-    // total number of active sessions
-    private int activeSessions;
-
     // pool of available object sessions
     private final Pool objSessions;
 
@@ -116,6 +115,11 @@ final class SessionManager {
         return (maxSessions <= DEFAULT_MAX_SESSIONS);
     }
 
+    // returns the total number of active sessions
+    int totalSessionCount() {
+        return SessionRef.totalCount();
+    }
+
     synchronized Session getObjSession() throws PKCS11Exception {
         Session session = objSessions.poll();
         if (session != null) {
@@ -136,7 +140,8 @@ final class SessionManager {
         }
         // create a new session rather than re-using an obj session
         // that avoids potential expensive cancels() for Signatures & RSACipher
-        if (activeSessions < maxSessions) {
+        if (maxSessions == Integer.MAX_VALUE ||
+                totalSessionCount() < maxSessions) {
             session = openSession();
             return ensureValid(session);
         }
@@ -159,14 +164,10 @@ final class SessionManager {
         if (debug != null) {
             String location = new Exception().getStackTrace()[2].toString();
             System.out.println("Killing session (" + location + ") active: "
-                + activeSessions);
+                + totalSessionCount());
         }
-        try {
         closeSession(session);
         return null;
-        } catch (PKCS11Exception e) {
-            throw new ProviderException(e);
-        }
     }
 
     synchronized Session releaseSession(Session session) {
@@ -187,7 +188,8 @@ final class SessionManager {
             return;
         }
         if (debug != null) {
-            System.out.println("Demoting session, active: " + activeSessions);
+            System.out.println("Demoting session, active: " +
+                totalSessionCount());
         }
         boolean present = objSessions.remove(session);
         if (present == false) {
@@ -199,16 +201,17 @@ final class SessionManager {
     }
 
     private Session openSession() throws PKCS11Exception {
-        if (activeSessions >= maxSessions) {
+        if ((maxSessions != Integer.MAX_VALUE) &&
+                (totalSessionCount() >= maxSessions)) {
             throw new ProviderException("No more sessions available");
         }
         long id = token.p11.C_OpenSession
                     (token.provider.slotID, openSessionFlags, null, null);
         Session session = new Session(token, id);
-        activeSessions++;
         if (debug != null) {
-            if (activeSessions > maxActiveSessions) {
-                maxActiveSessions = activeSessions;
+            int currTotal = totalSessionCount();
+            if (currTotal > maxActiveSessions) {
+                maxActiveSessions = currTotal;
                 if (maxActiveSessions % 10 == 0) {
                     System.out.println("Open sessions: " + maxActiveSessions);
                 }
@@ -217,13 +220,8 @@ final class SessionManager {
         return session;
     }
 
-    private void closeSession(Session session) throws PKCS11Exception {
-        if (session.hasObjects()) {
-            throw new ProviderException
-                ("Internal error: close session with active objects");
-        }
-        token.p11.C_CloseSession(session.id());
-        activeSessions--;
+    private void closeSession(Session session) {
+        session.close();
     }
 
     private static final class Pool {
@@ -267,28 +265,20 @@ final class SessionManager {
             }
             Collections.sort(pool);
             int i = 0;
-            PKCS11Exception exc = null;
             while (i < n - 1) { // always keep at least 1 session open
                 oldestSession = pool.get(i);
                 if (oldestSession.isLive(time)) {
                     break;
                 }
                 i++;
-                try {
                 mgr.closeSession(oldestSession);
-                } catch (PKCS11Exception e) {
-                    exc = e;
-                }
             }
             if (debug != null) {
                 System.out.println("Closing " + i + " idle sessions, active: "
-                        + mgr.activeSessions);
+                        + mgr.totalSessionCount());
             }
             List<Session> subList = pool.subList(0, i);
             subList.clear();
-            if (exc != null) {
-                throw new ProviderException(exc);
-            }
         }
 
     }