conn: fix race evicting statements on ErrUnPrepared (#1391)

Ensure we only remove prepared statements on ErrUnPrepared when our
cached statements ID matches the unprepared statement.

conn.go

@@ -1193,7 +1193,7 @@ func (c *Conn) executeQuery(ctx context.Context, qry *Query) *Iter {
 		return iter
 	case *RequestErrUnprepared:
 		stmtCacheKey := c.session.stmtsLRU.keyFor(c.addr, c.currentKeyspace, qry.stmt)
-		c.session.stmtsLRU.remove(stmtCacheKey)
+		c.session.stmtsLRU.evictPreparedID(stmtCacheKey, x.StatementId)
 		return c.executeQuery(ctx, qry)
 	case error:
 		return &Iter{err: x, framer: framer}
@@ -1334,7 +1334,7 @@ func (c *Conn) executeBatch(ctx context.Context, batch *Batch) *Iter {
 		stmt, found := stmts[string(x.StatementId)]
 		if found {
 			key := c.session.stmtsLRU.keyFor(c.addr, c.currentKeyspace, stmt)
-			c.session.stmtsLRU.remove(key)
+			c.session.stmtsLRU.evictPreparedID(key, x.StatementId)
 		}
 		return c.executeBatch(ctx, batch)
 	case *resultRowsFrame:

prepared_cache.go

@@ -1,6 +1,7 @@
 package gocql
 
 import (
+	"bytes"
 	"github.com/gocql/gocql/internal/lru"
 	"sync"
 )
@@ -59,6 +60,30 @@ func (p *preparedLRU) execIfMissing(key string, fn func(lru *lru.Cache) *infligh
 }
 
 func (p *preparedLRU) keyFor(addr, keyspace, statement string) string {
-	// TODO: maybe use []byte for keys?
+	// TODO: we should just use a struct for the key in the map
 	return addr + keyspace + statement
 }
+
+func (p *preparedLRU) evictPreparedID(key string, id []byte) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	val, ok := p.lru.Get(key)
+	if !ok {
+		return
+	}
+
+	ifp, ok := val.(*inflightPrepare)
+	if !ok {
+		return
+	}
+
+	select {
+	case <-ifp.done:
+		if bytes.Equal(id, ifp.preparedStatment.id) {
+			p.lru.Remove(key)
+		}
+	default:
+	}
+
+}