Repository: gocql/gocql
Pull request #1606 (2022-03-18 10:29:37)

Stability improvements
These are changes in kiwicom/gocql that are upstreamable. The first commit was submitted separately as #1605, but since some of the other commits depend on it, I included it here as well instead of rebasing the changes (and then needing to rebase back once the other PR is merged).

These are mainly improvements for connections in case of network issues.


Patch:
diff --git a/cluster.go b/cluster.go
index e7c4669..e700f40 100644
--- a/cluster.go
+++ b/cluster.go
@@ -48,6 +48,7 @@ type ClusterConfig struct {
 	ProtoVersion       int
 	Timeout            time.Duration                            // connection timeout (default: 600ms)
 	ConnectTimeout     time.Duration                            // initial connection timeout, used during initial dial to server (default: 600ms)
+	WriteTimeout       time.Duration                            // timeout for writing a query. defaults to Timeout if not specified.
 	Port               int                                      // port (default: 9042)
 	Keyspace           string                                   // initial keyspace (optional)
 	NumConns           int                                      // number of connections per host (default: 2)
@@ -135,6 +136,10 @@ type ClusterConfig struct {
 	// Use it to collect metrics / stats from frames by providing an implementation of FrameHeaderObserver.
 	FrameHeaderObserver FrameHeaderObserver
 
+	// StreamObserver will be notified of stream state changes.
+	// This can be used to track in-flight protocol requests and responses.
+	StreamObserver StreamObserver
+
 	// Default idempotence for queries
 	DefaultIdempotence bool
 
diff --git a/conn.go b/conn.go
index ccafca7..a3b4d41 100644
--- a/conn.go
+++ b/conn.go
@@ -120,6 +120,7 @@ type ConnConfig struct {
 	ProtoVersion   int
 	CQLVersion     string
 	Timeout        time.Duration
+	WriteTimeout   time.Duration
 	ConnectTimeout time.Duration
 	Dialer         Dialer
 	Compressor     Compressor
@@ -163,17 +164,22 @@ var TimeoutLimit int64 = 0
 type Conn struct {
 	conn net.Conn
 	r    *bufio.Reader
-	w    io.Writer
+	w    contextWriter
 
-	timeout       time.Duration
-	cfg           *ConnConfig
-	frameObserver FrameHeaderObserver
+	timeout        time.Duration
+	writeTimeout   time.Duration
+	cfg            *ConnConfig
+	frameObserver  FrameHeaderObserver
+	streamObserver StreamObserver
 
 	headerBuf [maxFrameHeaderSize]byte
 
 	streams *streams.IDGenerator
 	mu      sync.Mutex
-	calls   map[int]*callReq
+	// calls stores a map from stream ID to callReq.
+	// This map is protected by mu.
+	// calls should not be used when closed is true, calls is set to nil when closed=true.
+	calls map[int]*callReq
 
 	errorHandler ConnErrorHandler
 	compressor   Compressor
@@ -186,7 +192,9 @@ type Conn struct {
 
 	session *Session
 
-	closed int32
+	// true if connection close process for the connection started.
+	// closed is protected by mu.
+	closed bool
 	ctx    context.Context
 	cancel context.CancelFunc
 
@@ -271,6 +279,11 @@ func (s *Session) dialWithoutObserver(ctx context.Context, host *HostInfo, cfg *
 		conn = tconn
 	}
 
+	writeTimeout := cfg.Timeout
+	if cfg.WriteTimeout > 0 {
+		writeTimeout = cfg.WriteTimeout
+	}
+
 	ctx, cancel := context.WithCancel(ctx)
 	c := &Conn{
 		conn:          conn,
@@ -285,13 +298,17 @@ func (s *Session) dialWithoutObserver(ctx context.Context, host *HostInfo, cfg *
 		streams:       streams.New(cfg.ProtoVersion),
 		host:          host,
 		frameObserver: s.frameObserver,
-		w: &deadlineWriter{
-			w:       conn,
-			timeout: cfg.Timeout,
+		w: &deadlineContextWriter{
+			w:         conn,
+			timeout:   writeTimeout,
+			semaphore: make(chan struct{}, 1),
+			quit:      make(chan struct{}),
 		},
-		ctx:    ctx,
-		cancel: cancel,
-		logger: cfg.logger(),
+		ctx:            ctx,
+		cancel:         cancel,
+		logger:         cfg.logger(),
+		streamObserver: s.streamObserver,
+		writeTimeout:   writeTimeout,
 	}
 
 	if err := c.init(ctx); err != nil {
@@ -328,7 +345,7 @@ func (c *Conn) init(ctx context.Context) error {
 
 	// dont coalesce startup frames
 	if c.session.cfg.WriteCoalesceWaitTime > 0 && !c.cfg.disableCoalesce {
-		c.w = newWriteCoalescer(c.conn, c.timeout, c.session.cfg.WriteCoalesceWaitTime, ctx.Done())
+		c.w = newWriteCoalescer(c.conn, c.writeTimeout, c.session.cfg.WriteCoalesceWaitTime, ctx.Done())
 	}
 
 	go c.serve(ctx)
@@ -338,7 +355,7 @@ func (c *Conn) init(ctx context.Context) error {
 }
 
 func (c *Conn) Write(p []byte) (n int, err error) {
-	return c.w.Write(p)
+	return c.w.writeContext(context.Background(), p)
 }
 
 func (c *Conn) Read(p []byte) (n int, err error) {
@@ -414,7 +431,7 @@ func (s *startupCoordinator) setupConn(ctx context.Context) error {
 	return nil
 }
 
-func (s *startupCoordinator) write(ctx context.Context, frame frameWriter) (frame, error) {
+func (s *startupCoordinator) write(ctx context.Context, frame frameBuilder) (frame, error) {
 	select {
 	case s.frameTicker <- struct{}{}:
 	case <-ctx.Done():
@@ -521,23 +538,42 @@ func (s *startupCoordinator) authenticateHandshake(ctx context.Context, authFram
 }
 
 func (c *Conn) closeWithError(err error) {
-	if !atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
+	if c == nil {
 		return
 	}
 
-	// we should attempt to deliver the error back to the caller if it
-	// exists
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return
+	}
+	c.closed = true
+
+	var callsToClose map[int]*callReq
+
+	// We should attempt to deliver the error back to the caller if it
+	// exists. However, don't block c.mu while we are delivering the
+	// error to outstanding calls.
 	if err != nil {
-		c.mu.Lock()
-		for _, req := range c.calls {
-			// we need to send the error to all waiting queries, put the state
-			// of this conn into not active so that it can not execute any queries.
-			select {
-			case req.resp <- err:
-			case <-req.timeout:
-			}
+		callsToClose = c.calls
+		// It is safe to change c.calls to nil. Nobody should use it after c.closed is set to true.
+		c.calls = nil
+	}
+	c.mu.Unlock()
+
+	for _, req := range callsToClose {
+		// we need to send the error to all waiting queries.
+		select {
+		case req.resp <- callResp{err: err}:
+		case <-req.timeout:
+		}
+		if req.streamObserverContext != nil {
+			req.streamObserverEndOnce.Do(func() {
+				req.streamObserverContext.StreamAbandoned(ObservedStream{
+					Host: c.host,
+				})
+			})
 		}
-		c.mu.Unlock()
 	}
 
 	// if error was nil then unblock the quit channel
@@ -672,8 +708,8 @@ func (c *Conn) recv(ctx context.Context) error {
 		return fmt.Errorf("gocql: frame header stream is beyond call expected bounds: %d", head.stream)
 	} else if head.stream == -1 {
 		// TODO: handle cassandra event frames, we shouldnt get any currently
-		framer := newFramer(c, c, c.compressor, c.version)
-		if err := framer.readFrame(&head); err != nil {
+		framer := newFramer(c.compressor, c.version)
+		if err := framer.readFrame(c, &head); err != nil {
 			return err
 		}
 		go c.session.handleEvent(framer)
@@ -681,8 +717,8 @@ func (c *Conn) recv(ctx context.Context) error {
 	} else if head.stream <= 0 {
 		// reserved stream that we dont use, probably due to a protocol error
 		// or a bug in Cassandra, this should be an error, parse it and return.
-		framer := newFramer(c, c, c.compressor, c.version)
-		if err := framer.readFrame(&head); err != nil {
+		framer := newFramer(c.compressor, c.version)
+		if err := framer.readFrame(c, &head); err != nil {
 			return err
 		}
 
@@ -697,17 +733,23 @@ func (c *Conn) recv(ctx context.Context) error {
 	}
 
 	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return ErrConnectionClosed
+	}
 	call, ok := c.calls[head.stream]
 	delete(c.calls, head.stream)
 	c.mu.Unlock()
-	if call == nil || call.framer == nil || !ok {
+	if call == nil || !ok {
 		c.logger.Printf("gocql: received response for stream which has no handler: header=%v\n", head)
 		return c.discardFrame(head)
 	} else if head.stream != call.streamID {
 		panic(fmt.Sprintf("call has incorrect streamID: got %d expected %d", call.streamID, head.stream))
 	}
 
-	err = call.framer.readFrame(&head)
+	framer := newFramer(c.compressor, c.version)
+
+	err = framer.readFrame(c, &head)
 	if err != nil {
 		// only net errors should cause the connection to be closed. Though
 		// cassandra returning corrupt frames will be returned here as well.
@@ -719,7 +761,7 @@ func (c *Conn) recv(ctx context.Context) error {
 	// we either, return a response to the caller, the caller timedout, or the
 	// connection has closed. Either way we should never block indefinatly here
 	select {
-	case call.resp <- err:
+	case call.resp <- callResp{framer: framer, err: err}:
 	case <-call.timeout:
 		c.releaseStream(call)
 	case <-ctx.Done():
@@ -734,6 +776,14 @@ func (c *Conn) releaseStream(call *callReq) {
 	}
 
 	c.streams.Clear(call.streamID)
+
+	if call.streamObserverContext != nil {
+		call.streamObserverEndOnce.Do(func() {
+			call.streamObserverContext.StreamFinished(ObservedStream{
+				Host: c.host,
+			})
+		})
+	}
 }
 
 func (c *Conn) handleTimeout() {
@@ -743,152 +793,255 @@ func (c *Conn) handleTimeout() {
 }
 
 type callReq struct {
-	// could use a waitgroup but this allows us to do timeouts on the read/send
-	resp     chan error
-	framer   *framer
-	timeout  chan struct{} // indicates to recv() that a call has timedout
+	// resp will receive the frame that was sent as a response to this stream.
+	resp     chan callResp
+	timeout  chan struct{} // indicates to recv() that a call has timed out
 	streamID int           // current stream in use
 
 	timer *time.Timer
+
+	// streamObserverContext is notified about events regarding this stream
+	streamObserverContext StreamObserverContext
+
+	// streamObserverEndOnce ensures that either StreamAbandoned or StreamFinished is called,
+	// but not both.
+	streamObserverEndOnce sync.Once
 }
 
-type deadlineWriter struct {
-	w interface {
-		SetWriteDeadline(time.Time) error
-		io.Writer
-	}
+type callResp struct {
+	// framer is the response frame.
+	// May be nil if err is not nil.
+	framer *framer
+	// err is error encountered, if any.
+	err error
+}
+
+// contextWriter is like io.Writer, but takes context as well.
+type contextWriter interface {
+	// writeContext writes p to the connection.
+	//
+	// If ctx is canceled before we start writing p (e.g. during waiting while another write is currently in progress),
+	// p is not written and ctx.Err() is returned. Context is ignored after we start writing p (i.e. we don't interrupt
+	// blocked writes that are in progress) so that we always either write the full frame or not write it at all.
+	//
+	// It returns the number of bytes written from p (0 <= n <= len(p)) and any error that caused the write to stop
+	// early. writeContext must return a non-nil error if it returns n < len(p). writeContext must not modify the
+	// data in p, even temporarily.
+	writeContext(ctx context.Context, p []byte) (n int, err error)
+}
+
+type deadlineWriter interface {
+	SetWriteDeadline(time.Time) error
+	io.Writer
+}
+
+type deadlineContextWriter struct {
+	w       deadlineWriter
 	timeout time.Duration
+	// semaphore protects critical section for SetWriteDeadline/Write.
+	// It is a channel with capacity 1.
+	semaphore chan struct{}
+
+	// quit closed once the connection is closed.
+	quit chan struct{}
 }
 
-func (c *deadlineWriter) Write(p []byte) (int, error) {
+// writeContext implements contextWriter.
+func (c *deadlineContextWriter) writeContext(ctx context.Context, p []byte) (int, error) {
+	select {
+	case <-ctx.Done():
+		return 0, ctx.Err()
+	case <-c.quit:
+		return 0, ErrConnectionClosed
+	case c.semaphore <- struct{}{}:
+		// acquired
+	}
+
+	defer func() {
+		// release
+		<-c.semaphore
+	}()
+
 	if c.timeout > 0 {
-		c.w.SetWriteDeadline(time.Now().Add(c.timeout))
+		err := c.w.SetWriteDeadline(time.Now().Add(c.timeout))
+		if err != nil {
+			return 0, err
+		}
 	}
 	return c.w.Write(p)
 }
 
-func newWriteCoalescer(conn net.Conn, timeout time.Duration, d time.Duration, quit <-chan struct{}) *writeCoalescer {
+func newWriteCoalescer(conn deadlineWriter, writeTimeout, coalesceDuration time.Duration,
+	quit <-chan struct{}) *writeCoalescer {
 	wc := &writeCoalescer{
-		writeCh: make(chan struct{}), // TODO: could this be sync?
-		cond:    sync.NewCond(&sync.Mutex{}),
+		writeCh: make(chan writeRequest),
 		c:       conn,
 		quit:    quit,
-		timeout: timeout,
+		timeout: writeTimeout,
 	}
-	go wc.writeFlusher(d)
+	go wc.writeFlusher(coalesceDuration)
 	return wc
 }
 
 type writeCoalescer struct {
-	c net.Conn
+	c deadlineWriter
+
+	mu sync.Mutex
 
 	quit    <-chan struct{}
-	writeCh chan struct{}
-	running bool
+	writeCh chan writeRequest
 
-	// cond waits for the buffer to be flushed
-	cond    *sync.Cond
-	buffers net.Buffers
 	timeout time.Duration
 
-	// result of the write
-	err error
-}
-
-func (w *writeCoalescer) flushLocked() {
-	w.running = false
-	if len(w.buffers) == 0 {
-		return
-	}
-
-	if w.timeout > 0 {
-		w.c.SetWriteDeadline(time.Now().Add(w.timeout))
-	}
-
-	// Given we are going to do a fanout n is useless and according to
-	// the docs WriteTo should return 0 and err or bytes written and
-	// no error.
-	_, w.err = w.buffers.WriteTo(w.c)
-	if w.err != nil {
-		w.buffers = nil
-	}
-	w.cond.Broadcast()
+	testEnqueuedHook func()
+	testFlushedHook  func()
 }
 
-func (w *writeCoalescer) flush() {
-	w.cond.L.Lock()
-	w.flushLocked()
-	w.cond.L.Unlock()
+type writeRequest struct {
+	// resultChan is a channel (with buffer size 1) where to send results of the write.
+	resultChan chan<- writeResult
+	// data to write.
+	data []byte
 }
 
-func (w *writeCoalescer) stop() {
-	w.cond.L.Lock()
-	defer w.cond.L.Unlock()
-
-	w.flushLocked()
-	// nil the channel out sends block forever on it
-	// instead of closing which causes a send on closed channel
-	// panic.
-	w.writeCh = nil
+type writeResult struct {
+	n   int
+	err error
 }
 
-func (w *writeCoalescer) Write(p []byte) (int, error) {
-	w.cond.L.Lock()
-
-	if !w.running {
-		select {
-		case w.writeCh <- struct{}{}:
-			w.running = true
-		case <-w.quit:
-			w.cond.L.Unlock()
-			return 0, io.EOF // TODO: better error here?
-		}
+// writeContext implements contextWriter.
+func (w *writeCoalescer) writeContext(ctx context.Context, p []byte) (int, error) {
+	resultChan := make(chan writeResult, 1)
+	wr := writeRequest{
+		resultChan: resultChan,
+		data:       p,
 	}
 
-	w.buffers = append(w.buffers, p)
-	for len(w.buffers) != 0 {
-		w.cond.Wait()
+	select {
+	case <-ctx.Done():
+		return 0, ctx.Err()
+	case <-w.quit:
+		return 0, io.EOF // TODO: better error here?
+	case w.writeCh <- wr:
+		// enqueued for writing
 	}
 
-	err := w.err
-	w.cond.L.Unlock()
-
-	if err != nil {
-		return 0, err
+	if w.testEnqueuedHook != nil {
+		w.testEnqueuedHook()
 	}
-	return len(p), nil
+
+	result := <-resultChan
+	return result.n, result.err
 }
 
 func (w *writeCoalescer) writeFlusher(interval time.Duration) {
 	timer := time.NewTimer(interval)
 	defer timer.Stop()
-	defer w.stop()
 
 	if !timer.Stop() {
 		<-timer.C
 	}
 
+	w.writeFlusherImpl(timer.C, func() { timer.Reset(interval) })
+}
+
+func (w *writeCoalescer) writeFlusherImpl(timerC <-chan time.Time, resetTimer func()) {
+	running := false
+
+	var buffers net.Buffers
+	var resultChans []chan<- writeResult
+
 	for {
-		// wait for a write to start the flush loop
 		select {
-		case <-w.writeCh:
+		case req := <-w.writeCh:
+			buffers = append(buffers, req.data)
+			resultChans = append(resultChans, req.resultChan)
+			if !running {
+				// Start timer on first write.
+				resetTimer()
+				running = true
+			}
 		case <-w.quit:
+			result := writeResult{
+				n:   0,
+				err: io.EOF, // TODO: better error here?
+			}
+			// Unblock whoever was waiting.
+			for _, resultChan := range resultChans {
+				// resultChan has capacity 1, so it does not block.
+				resultChan <- result
+			}
 			return
+		case <-timerC:
+			running = false
+			w.flush(resultChans, buffers)
+			buffers = nil
+			resultChans = nil
+			if w.testFlushedHook != nil {
+				w.testFlushedHook()
+			}
 		}
+	}
+}
 
-		timer.Reset(interval)
-
-		select {
-		case <-w.quit:
+func (w *writeCoalescer) flush(resultChans []chan<- writeResult, buffers net.Buffers) {
+	// Flush everything we have so far.
+	if w.timeout > 0 {
+		err := w.c.SetWriteDeadline(time.Now().Add(w.timeout))
+		if err != nil {
+			for i := range resultChans {
+				resultChans[i] <- writeResult{
+					n:   0,
+					err: err,
+				}
+			}
 			return
-		case <-timer.C:
 		}
+	}
+	// Copy buffers because WriteTo modifies buffers in-place.
+	buffers2 := make(net.Buffers, len(buffers))
+	copy(buffers2, buffers)
+	n, err := buffers2.WriteTo(w.c)
+	// Writes of bytes before n succeeded, writes of bytes starting from n failed with err.
+	// Use n as remaining byte counter.
+	for i := range buffers {
+		if int64(len(buffers[i])) <= n {
+			// this buffer was fully written.
+			resultChans[i] <- writeResult{
+				n:   len(buffers[i]),
+				err: nil,
+			}
+			n -= int64(len(buffers[i]))
+		} else {
+			// this buffer was not (fully) written.
+			resultChans[i] <- writeResult{
+				n:   int(n),
+				err: err,
+			}
+			n = 0
+		}
+	}
+}
 
-		w.flush()
+// addCall attempts to add a call to c.calls.
+// It fails with error if the connection already started closing or if a call for the given stream
+// already exists.
+func (c *Conn) addCall(call *callReq) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.closed {
+		return ErrConnectionClosed
+	}
+	existingCall := c.calls[call.streamID]
+	if existingCall != nil {
+		return fmt.Errorf("attempting to use stream already in use: %d -> %d", call.streamID,
+			existingCall.streamID)
 	}
+	c.calls[call.streamID] = call
+	return nil
 }
 
-func (c *Conn) exec(ctx context.Context, req frameWriter, tracer Tracer) (*framer, error) {
+func (c *Conn) exec(ctx context.Context, req frameBuilder, tracer Tracer) (*framer, error) {
 	if ctxErr := ctx.Err(); ctxErr != nil {
 		return nil, ctxErr
 	}
@@ -900,42 +1053,79 @@ func (c *Conn) exec(ctx context.Context, req frameWriter, tracer Tracer) (*frame
 	}
 
 	// resp is basically a waiting semaphore protecting the framer
-	framer := newFramer(c, c, c.compressor, c.version)
+	framer := newFramer(c.compressor, c.version)
 
 	call := &callReq{
-		framer:   framer,
 		timeout:  make(chan struct{}),
 		streamID: stream,
-		resp:     make(chan error),
+		resp:     make(chan callResp),
 	}
 
-	c.mu.Lock()
-	existingCall := c.calls[stream]
-	if existingCall == nil {
-		c.calls[stream] = call
+	if c.streamObserver != nil {
+		call.streamObserverContext = c.streamObserver.StreamContext(ctx)
 	}
-	c.mu.Unlock()
 
-	if existingCall != nil {
-		return nil, fmt.Errorf("attempting to use stream already in use: %d -> %d", stream, existingCall.streamID)
+	if err := c.addCall(call); err != nil {
+		return nil, err
 	}
 
+	// After this point, we need to either read from call.resp or close(call.timeout)
+	// since closeWithError can try to write a connection close error to call.resp.
+	// If we don't close(call.timeout) or read from call.resp, closeWithError can deadlock.
+
 	if tracer != nil {
 		framer.trace()
 	}
 
-	err := req.writeFrame(framer, stream)
+	if call.streamObserverContext != nil {
+		call.streamObserverContext.StreamStarted(ObservedStream{
+			Host: c.host,
+		})
+	}
+
+	err := req.buildFrame(framer, stream)
+	if err != nil {
+		// closeWithError will block waiting for this stream to either receive a response
+		// or for us to timeout.
+		close(call.timeout)
+		// We failed to serialize the frame into a buffer.
+		// This should not affect the connection as we didn't write anything. We just free the current call.
+		c.mu.Lock()
+		if !c.closed {
+			delete(c.calls, call.streamID)
+		}
+		c.mu.Unlock()
+		// We need to release the stream after we remove the call from c.calls, otherwise the existingCall != nil
+		// check above could fail.
+		c.releaseStream(call)
+		return nil, err
+	}
+
+	n, err := c.w.writeContext(ctx, framer.buf)
 	if err != nil {
 		// closeWithError will block waiting for this stream to either receive a response
 		// or for us to timeout, close the timeout chan here. Im not entirely sure
 		// but we should not get a response after an error on the write side.
 		close(call.timeout)
-		// I think this is the correct thing to do, im not entirely sure. It is not
-		// ideal as readers might still get some data, but they probably wont.
-		// Here we need to be careful as the stream is not available and if all
-		// writes just timeout or fail then the pool might use this connection to
-		// send a frame on, with all the streams used up and not returned.
-		c.closeWithError(err)
+		if (errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) && n == 0 {
+			// We have not started to write this frame.
+			// Release the stream as no response can come from the server on the stream.
+			c.mu.Lock()
+			if !c.closed {
+				delete(c.calls, call.streamID)
+			}
+			c.mu.Unlock()
+			// We need to release the stream after we remove the call from c.calls, otherwise the existingCall != nil
+			// check above could fail.
+			c.releaseStream(call)
+		} else {
+			// I think this is the correct thing to do, im not entirely sure. It is not
+			// ideal as readers might still get some data, but they probably wont.
+			// Here we need to be careful as the stream is not available and if all
+			// writes just timeout or fail then the pool might use this connection to
+			// send a frame on, with all the streams used up and not returned.
+			c.closeWithError(err)
+		}
 		return nil, err
 	}
 
@@ -963,9 +1153,9 @@ func (c *Conn) exec(ctx context.Context, req frameWriter, tracer Tracer) (*frame
 	}
 
 	select {
-	case err := <-call.resp:
+	case resp := <-call.resp:
 		close(call.timeout)
-		if err != nil {
+		if resp.err != nil {
 			if !c.Closed() {
 				// if the connection is closed then we cant release the stream,
 				// this is because the request is still outstanding and we have
@@ -973,8 +1163,21 @@ func (c *Conn) exec(ctx context.Context, req frameWriter, tracer Tracer) (*frame
 				// connection to close.
 				c.releaseStream(call)
 			}
-			return nil, err
+			return nil, resp.err
+		}
+		// dont release the stream if detect a timeout as another request can reuse
+		// that stream and get a response for the old request, which we have no
+		// easy way of detecting.
+		//
+		// Ensure that the stream is not released if there are potentially outstanding
+		// requests on the stream to prevent nil pointer dereferences in recv().
+		defer c.releaseStream(call)
+
+		if v := resp.framer.header.version.version(); v != c.version {
+			return nil, NewErrProtocol("unexpected protocol version in response: got %d expected %d", v, c.version)
 		}
+
+		return resp.framer, nil
 	case <-timeoutCh:
 		close(call.timeout)
 		c.handleTimeout()
@@ -983,22 +1186,55 @@ func (c *Conn) exec(ctx context.Context, req frameWriter, tracer Tracer) (*frame
 		close(call.timeout)
 		return nil, ctx.Err()
 	case <-c.ctx.Done():
+		close(call.timeout)
 		return nil, ErrConnectionClosed
 	}
+}
 
-	// dont release the stream if detect a timeout as another request can reuse
-	// that stream and get a response for the old request, which we have no
-	// easy way of detecting.
-	//
-	// Ensure that the stream is not released if there are potentially outstanding
-	// requests on the stream to prevent nil pointer dereferences in recv().
-	defer c.releaseStream(call)
+// ObservedStream observes a single request/response stream.
+type ObservedStream struct {
+	// Host of the connection used to send the stream.
+	Host *HostInfo
+}
 
-	if v := framer.header.version.version(); v != c.version {
-		return nil, NewErrProtocol("unexpected protocol version in response: got %d expected %d", v, c.version)
-	}
+// StreamObserver is notified about request/response pairs.
+// Streams are created for executing queries/batches or
+// internal requests to the database and might live longer than
+// execution of the query - the stream is still tracked until
+// response arrives so that stream IDs are not reused.
+type StreamObserver interface {
+	// StreamContext is called before creating a new stream.
+	// ctx is context passed to Session.Query / Session.Batch,
+	// but might also be an internal context (for example
+	// for internal requests that use control connection).
+	// StreamContext might return nil if it is not interested
+	// in the details of this stream.
+	// StreamContext is called before the stream is created
+	// and the returned StreamObserverContext might be discarded
+	// without any methods called on the StreamObserverContext if
+	// creation of the stream fails.
+	// Note that if you don't need to track per-stream data,
+	// you can always return the same StreamObserverContext.
+	StreamContext(ctx context.Context) StreamObserverContext
+}
 
-	return framer, nil
+// StreamObserverContext is notified about state of a stream.
+// A stream is started every time a request is written to the server
+// and is finished when a response is received.
+// It is abandoned when the underlying network connection is closed
+// before receiving a response.
+type StreamObserverContext interface {
+	// StreamStarted is called when the stream is started.
+	// This happens just before a request is written to the wire.
+	StreamStarted(observedStream ObservedStream)
+
+	// StreamAbandoned is called when we stop waiting for response.
+	// This happens when the underlying network connection is closed.
+	// StreamFinished won't be called if StreamAbandoned is.
+	StreamAbandoned(observedStream ObservedStream)
+
+	// StreamFinished is called when we receive a response for the stream.
+	StreamFinished(observedStream ObservedStream)
 }
 
 type preparedStatment struct {
@@ -1130,7 +1366,7 @@ func (c *Conn) executeQuery(ctx context.Context, qry *Query) *Iter {
 	}
 
 	var (
-		frame frameWriter
+		frame frameBuilder
 		info  *preparedStatment
 	)
 
@@ -1271,7 +1507,9 @@ func (c *Conn) Pick(qry *Query) *Conn {
 }
 
 func (c *Conn) Closed() bool {
-	return atomic.LoadInt32(&c.closed) == 1
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.closed
 }
 
 func (c *Conn) Address() string {
diff --git a/connectionpool.go b/connectionpool.go
index e71a376..d4c1521 100644
--- a/connectionpool.go
+++ b/connectionpool.go
@@ -107,6 +107,7 @@ func connConfig(cfg *ClusterConfig) (*ConnConfig, error) {
 		ProtoVersion:    cfg.ProtoVersion,
 		CQLVersion:      cfg.CQLVersion,
 		Timeout:         cfg.Timeout,
+		WriteTimeout:    cfg.WriteTimeout,
 		ConnectTimeout:  cfg.ConnectTimeout,
 		Dialer:          cfg.Dialer,
 		Compressor:      cfg.Compressor,
diff --git a/control.go b/control.go
index 0501366..8bc5310 100644
--- a/control.go
+++ b/control.go
@@ -419,7 +419,7 @@ func (c *controlConn) getConn() *connHost {
 	return c.conn.Load().(*connHost)
 }
 
-func (c *controlConn) writeFrame(w frameWriter) (frame, error) {
+func (c *controlConn) writeFrame(w frameBuilder) (frame, error) {
 	ch := c.getConn()
 	if ch == nil {
 		return nil, errNoControl
diff --git a/frame.go b/frame.go
index 3a14284..fff55a4 100644
--- a/frame.go
+++ b/frame.go
@@ -364,9 +364,6 @@ type FrameHeaderObserver interface {
 
 // a framer is responsible for reading, writing and parsing frames on a single stream
 type framer struct {
-	r io.Reader
-	w io.Writer
-
 	proto byte
 	// flags are for outgoing flags, enabling compression and tracing etc
 	flags    byte
@@ -378,20 +375,20 @@ type framer struct {
 	// if tracing flag is set this is not nil
 	traceID []byte
 
-	// holds a ref to the whole byte slice for rbuf so that it can be reset to
+	// holds a ref to the whole byte slice for buf so that it can be reset to
 	// 0 after a read.
 	readBuffer []byte
 
-	rbuf []byte
-	wbuf []byte
+	buf []byte
 
 	customPayload map[string][]byte
 }
 
-func newFramer(r io.Reader, w io.Writer, compressor Compressor, version byte) *framer {
+func newFramer(compressor Compressor, version byte) *framer {
+	buf := make([]byte, defaultBufSize)
 	f := &framer{
-		wbuf:       make([]byte, defaultBufSize),
-		readBuffer: make([]byte, defaultBufSize),
+		buf:        buf[:0],
+		readBuffer: buf,
 	}
 	var flags byte
 	if compressor != nil {
@@ -413,12 +410,6 @@ func newFramer(r io.Reader, w io.Writer, compressor Compressor, version byte) *f
 	f.flags = flags
 	f.headSize = headSize
 
-	f.r = r
-	f.rbuf = f.readBuffer[:0]
-
-	f.w = w
-	f.wbuf = f.wbuf[:0]
-
 	f.header = nil
 	f.traceID = nil
 
@@ -488,12 +479,12 @@ func (f *framer) payload() {
 }
 
 // reads a frame form the wire into the framers buffer
-func (f *framer) readFrame(head *frameHeader) error {
+func (f *framer) readFrame(r io.Reader, head *frameHeader) error {
 	if head.length < 0 {
 		return fmt.Errorf("frame body length can not be less than 0: %d", head.length)
 	} else if head.length > maxFrameSize {
 		// need to free up the connection to be used again
-		_, err := io.CopyN(ioutil.Discard, f.r, int64(head.length))
+		_, err := io.CopyN(ioutil.Discard, r, int64(head.length))
 		if err != nil {
 			return fmt.Errorf("error whilst trying to discard frame with invalid length: %v", err)
 		}
@@ -501,14 +492,14 @@ func (f *framer) readFrame(head *frameHeader) error {
 	}
 
 	if cap(f.readBuffer) >= head.length {
-		f.rbuf = f.readBuffer[:head.length]
+		f.buf = f.readBuffer[:head.length]
 	} else {
 		f.readBuffer = make([]byte, head.length)
-		f.rbuf = f.readBuffer
+		f.buf = f.readBuffer
 	}
 
 	// assume the underlying reader takes care of timeouts and retries
-	n, err := io.ReadFull(f.r, f.rbuf)
+	n, err := io.ReadFull(r, f.buf)
 	if err != nil {
 		return fmt.Errorf("unable to read frame body: read %d/%d bytes: %v", n, head.length, err)
 	}
@@ -518,7 +509,7 @@ func (f *framer) readFrame(head *frameHeader) error {
 			return NewErrProtocol("no compressor available with compressed frame body")
 		}
 
-		f.rbuf, err = f.compres.Decode(f.rbuf)
+		f.buf, err = f.compres.Decode(f.buf)
 		if err != nil {
 			return err
 		}
@@ -711,25 +702,25 @@ func (f *framer) readErrorMap() (errMap ErrorMap) {
 }
 
 func (f *framer) writeHeader(flags byte, op frameOp, stream int) {
-	f.wbuf = f.wbuf[:0]
-	f.wbuf = append(f.wbuf,
+	f.buf = f.buf[:0]
+	f.buf = append(f.buf,
 		f.proto,
 		flags,
 	)
 
 	if f.proto > protoVersion2 {
-		f.wbuf = append(f.wbuf,
+		f.buf = append(f.buf,
 			byte(stream>>8),
 			byte(stream),
 		)
 	} else {
-		f.wbuf = append(f.wbuf,
+		f.buf = append(f.buf,
 			byte(stream),
 		)
 	}
 
 	// pad out length
-	f.wbuf = append(f.wbuf,
+	f.buf = append(f.buf,
 		byte(op),
 		0,
 		0,
@@ -744,43 +735,43 @@ func (f *framer) setLength(length int) {
 		p = 5
 	}
 
-	f.wbuf[p+0] = byte(length >> 24)
-	f.wbuf[p+1] = byte(length >> 16)
-	f.wbuf[p+2] = byte(length >> 8)
-	f.wbuf[p+3] = byte(length)
+	f.buf[p+0] = byte(length >> 24)
+	f.buf[p+1] = byte(length >> 16)
+	f.buf[p+2] = byte(length >> 8)
+	f.buf[p+3] = byte(length)
 }
 
-func (f *framer) finishWrite() error {
-	if len(f.wbuf) > maxFrameSize {
+func (f *framer) finish() error {
+	if len(f.buf) > maxFrameSize {
 		// huge app frame, lets remove it so it doesn't bloat the heap
-		f.wbuf = make([]byte, defaultBufSize)
+		f.buf = make([]byte, defaultBufSize)
 		return ErrFrameTooBig
 	}
 
-	if f.wbuf[1]&flagCompress == flagCompress {
+	if f.buf[1]&flagCompress == flagCompress {
 		if f.compres == nil {
 			panic("compress flag set with no compressor")
 		}
 
 		// TODO: only compress frames which are big enough
-		compressed, err := f.compres.Encode(f.wbuf[f.headSize:])
+		compressed, err := f.compres.Encode(f.buf[f.headSize:])
 		if err != nil {
 			return err
 		}
 
-		f.wbuf = append(f.wbuf[:f.headSize], compressed...)
+		f.buf = append(f.buf[:f.headSize], compressed...)
 	}
-	length := len(f.wbuf) - f.headSize
+	length := len(f.buf) - f.headSize
 	f.setLength(length)
 
-	_, err := f.w.Write(f.wbuf)
-	if err != nil {
-		return err
-	}
-
 	return nil
 }
 
+func (f *framer) writeTo(w io.Writer) error {
+	_, err := w.Write(f.buf)
+	return err
+}
+
 func (f *framer) readTrace() {
 	f.traceID = f.readUUID().Bytes()
 }
@@ -819,11 +810,11 @@ func (w writeStartupFrame) String() string {
 	return fmt.Sprintf("[startup opts=%+v]", w.opts)
 }
 
-func (w *writeStartupFrame) writeFrame(f *framer, streamID int) error {
+func (w *writeStartupFrame) buildFrame(f *framer, streamID int) error {
 	f.writeHeader(f.flags&^flagCompress, opStartup, streamID)
 	f.writeStringMap(w.opts)
 
-	return f.finishWrite()
+	return f.finish()
 }
 
 type writePrepareFrame struct {
@@ -832,7 +823,7 @@ type writePrepareFrame struct {
 	customPayload map[string][]byte
 }
 
-func (w *writePrepareFrame) writeFrame(f *framer, streamID int) error {
+func (w *writePrepareFrame) buildFrame(f *framer, streamID int) error {
 	if len(w.customPayload) > 0 {
 		f.payload()
 	}
@@ -855,7 +846,7 @@ func (w *writePrepareFrame) writeFrame(f *framer, streamID int) error {
 		f.writeString(w.keyspace)
 	}
 
-	return f.finishWrite()
+	return f.finish()
 }
 
 func (f *framer) readTypeInfo() TypeInfo {
@@ -1413,14 +1404,14 @@ func (a *writeAuthResponseFrame) String() string {
 	return fmt.Sprintf("[auth_response data=%q]", a.data)
 }
 
-func (a *writeAuthResponseFrame) writeFrame(framer *framer, streamID int) error {
+func (a *writeAuthResponseFrame) buildFrame(framer *framer, streamID int) error {
 	return framer.writeAuthResponseFrame(streamID, a.data)
 }
 
 func (f *framer) writeAuthResponseFrame(streamID int, data []byte) error {
 	f.writeHeader(f.flags, opAuthResponse, streamID)
 	f.writeBytes(data)
-	return f.finishWrite()
+	return f.finish()
 }
 
 type queryValues struct {
@@ -1558,7 +1549,7 @@ func (w *writeQueryFrame) String() string {
 	return fmt.Sprintf("[query statement=%q params=%v]", w.statement, w.params)
 }
 
-func (w *writeQueryFrame) writeFrame(framer *framer, streamID int) error {
+func (w *writeQueryFrame) buildFrame(framer *framer, streamID int) error {
 	return framer.writeQueryFrame(streamID, w.statement, &w.params, w.customPayload)
 }
 
@@ -1571,16 +1562,16 @@ func (f *framer) writeQueryFrame(streamID int, statement string, params *queryPa
 	f.writeLongString(statement)
 	f.writeQueryParams(params)
 
-	return f.finishWrite()
+	return f.finish()
 }
 
-type frameWriter interface {
-	writeFrame(framer *framer, streamID int) error
+type frameBuilder interface {
+	buildFrame(framer *framer, streamID int) error
 }
 
 type frameWriterFunc func(framer *framer, streamID int) error
 
-func (f frameWriterFunc) writeFrame(framer *framer, streamID int) error {
+func (f frameWriterFunc) buildFrame(framer *framer, streamID int) error {
 	return f(framer, streamID)
 }
 
@@ -1596,7 +1587,7 @@ func (e *writeExecuteFrame) String() string {
 	return fmt.Sprintf("[execute id=% X params=%v]", e.preparedID, &e.params)
 }
 
-func (e *writeExecuteFrame) writeFrame(fr *framer, streamID int) error {
+func (e *writeExecuteFrame) buildFrame(fr *framer, streamID int) error {
 	return fr.writeExecuteFrame(streamID, e.preparedID, &e.params, &e.customPayload)
 }
 
@@ -1622,7 +1613,7 @@ func (f *framer) writeExecuteFrame(streamID int, preparedID []byte, params *quer
 		f.writeConsistency(params.consistency)
 	}
 
-	return f.finishWrite()
+	return f.finish()
 }
 
 // TODO: can we replace BatchStatemt with batchStatement? As they prety much
@@ -1647,7 +1638,7 @@ type writeBatchFrame struct {
 	customPayload map[string][]byte
 }
 
-func (w *writeBatchFrame) writeFrame(framer *framer, streamID int) error {
+func (w *writeBatchFrame) buildFrame(framer *framer, streamID int) error {
 	return framer.writeBatchFrame(streamID, w, w.customPayload)
 }
 
@@ -1725,25 +1716,25 @@ func (f *framer) writeBatchFrame(streamID int, w *writeBatchFrame, customPayload
 		}
 	}
 
-	return f.finishWrite()
+	return f.finish()
 }
 
 type writeOptionsFrame struct{}
 
-func (w *writeOptionsFrame) writeFrame(framer *framer, streamID int) error {
+func (w *writeOptionsFrame) buildFrame(framer *framer, streamID int) error {
 	return framer.writeOptionsFrame(streamID, w)
 }
 
 func (f *framer) writeOptionsFrame(stream int, _ *writeOptionsFrame) error {
 	f.writeHeader(f.flags&^flagCompress, opOptions, stream)
-	return f.finishWrite()
+	return f.finish()
 }
 
 type writeRegisterFrame struct {
 	events []string
 }
 
-func (w *writeRegisterFrame) writeFrame(framer *framer, streamID int) error {
+func (w *writeRegisterFrame) buildFrame(framer *framer, streamID int) error {
 	return framer.writeRegisterFrame(streamID, w)
 }
 
@@ -1751,70 +1742,70 @@ func (f *framer) writeRegisterFrame(streamID int, w *writeRegisterFrame) error {
 	f.writeHeader(f.flags, opRegister, streamID)
 	f.writeStringList(w.events)
 
-	return f.finishWrite()
+	return f.finish()
 }
 
 func (f *framer) readByte() byte {
-	if len(f.rbuf) < 1 {
-		panic(fmt.Errorf("not enough bytes in buffer to read byte require 1 got: %d", len(f.rbuf)))
+	if len(f.buf) < 1 {
+		panic(fmt.Errorf("not enough bytes in buffer to read byte require 1 got: %d", len(f.buf)))
 	}
 
-	b := f.rbuf[0]
-	f.rbuf = f.rbuf[1:]
+	b := f.buf[0]
+	f.buf = f.buf[1:]
 	return b
 }
 
 func (f *framer) readInt() (n int) {
-	if len(f.rbuf) < 4 {
-		panic(fmt.Errorf("not enough bytes in buffer to read int require 4 got: %d", len(f.rbuf)))
+	if len(f.buf) < 4 {
+		panic(fmt.Errorf("not enough bytes in buffer to read int require 4 got: %d", len(f.buf)))
 	}
 
-	n = int(int32(f.rbuf[0])<<24 | int32(f.rbuf[1])<<16 | int32(f.rbuf[2])<<8 | int32(f.rbuf[3]))
-	f.rbuf = f.rbuf[4:]
+	n = int(int32(f.buf[0])<<24 | int32(f.buf[1])<<16 | int32(f.buf[2])<<8 | int32(f.buf[3]))
+	f.buf = f.buf[4:]
 	return
 }
 
 func (f *framer) readShort() (n uint16) {
-	if len(f.rbuf) < 2 {
-		panic(fmt.Errorf("not enough bytes in buffer to read short require 2 got: %d", len(f.rbuf)))
+	if len(f.buf) < 2 {
+		panic(fmt.Errorf("not enough bytes in buffer to read short require 2 got: %d", len(f.buf)))
 	}
-	n = uint16(f.rbuf[0])<<8 | uint16(f.rbuf[1])
-	f.rbuf = f.rbuf[2:]
+	n = uint16(f.buf[0])<<8 | uint16(f.buf[1])
+	f.buf = f.buf[2:]
 	return
 }
 
 func (f *framer) readString() (s string) {
 	size := f.readShort()
 
-	if len(f.rbuf) < int(size) {
-		panic(fmt.Errorf("not enough bytes in buffer to read string require %d got: %d", size, len(f.rbuf)))
+	if len(f.buf) < int(size) {
+		panic(fmt.Errorf("not enough bytes in buffer to read string require %d got: %d", size, len(f.buf)))
 	}
 
-	s = string(f.rbuf[:size])
-	f.rbuf = f.rbuf[size:]
+	s = string(f.buf[:size])
+	f.buf = f.buf[size:]
 	return
 }
 
 func (f *framer) readLongString() (s string) {
 	size := f.readInt()
 
-	if len(f.rbuf) < size {
-		panic(fmt.Errorf("not enough bytes in buffer to read long string require %d got: %d", size, len(f.rbuf)))
+	if len(f.buf) < size {
+		panic(fmt.Errorf("not enough bytes in buffer to read long string require %d got: %d", size, len(f.buf)))
 	}
 
-	s = string(f.rbuf[:size])
-	f.rbuf = f.rbuf[size:]
+	s = string(f.buf[:size])
+	f.buf = f.buf[size:]
 	return
 }
 
 func (f *framer) readUUID() *UUID {
-	if len(f.rbuf) < 16 {
-		panic(fmt.Errorf("not enough bytes in buffer to read uuid require %d got: %d", 16, len(f.rbuf)))
+	if len(f.buf) < 16 {
+		panic(fmt.Errorf("not enough bytes in buffer to read uuid require %d got: %d", 16, len(f.buf)))
 	}
 
 	// TODO: how to handle this error, if it is a uuid, then sureley, problems?
-	u, _ := UUIDFromBytes(f.rbuf[:16])
-	f.rbuf = f.rbuf[16:]
+	u, _ := UUIDFromBytes(f.buf[:16])
+	f.buf = f.buf[16:]
 	return &u
 }
 
@@ -1835,12 +1826,12 @@ func (f *framer) readBytesInternal() ([]byte, error) {
 		return nil, nil
 	}
 
-	if len(f.rbuf) < size {
-		return nil, fmt.Errorf("not enough bytes in buffer to read bytes require %d got: %d", size, len(f.rbuf))
+	if len(f.buf) < size {
+		return nil, fmt.Errorf("not enough bytes in buffer to read bytes require %d got: %d", size, len(f.buf))
 	}
 
-	l := f.rbuf[:size]
-	f.rbuf = f.rbuf[size:]
+	l := f.buf[:size]
+	f.buf = f.buf[size:]
 
 	return l, nil
 }
@@ -1856,35 +1847,35 @@ func (f *framer) readBytes() []byte {
 
 func (f *framer) readShortBytes() []byte {
 	size := f.readShort()
-	if len(f.rbuf) < int(size) {
-		panic(fmt.Errorf("not enough bytes in buffer to read short bytes: require %d got %d", size, len(f.rbuf)))
+	if len(f.buf) < int(size) {
+		panic(fmt.Errorf("not enough bytes in buffer to read short bytes: require %d got %d", size, len(f.buf)))
 	}
 
-	l := f.rbuf[:size]
-	f.rbuf = f.rbuf[size:]
+	l := f.buf[:size]
+	f.buf = f.buf[size:]
 
 	return l
 }
 
 func (f *framer) readInetAdressOnly() net.IP {
-	if len(f.rbuf) < 1 {
-		panic(fmt.Errorf("not enough bytes in buffer to read inet size require %d got: %d", 1, len(f.rbuf)))
+	if len(f.buf) < 1 {
+		panic(fmt.Errorf("not enough bytes in buffer to read inet size require %d got: %d", 1, len(f.buf)))
 	}
 
-	size := f.rbuf[0]
-	f.rbuf = f.rbuf[1:]
+	size := f.buf[0]
+	f.buf = f.buf[1:]
 
 	if !(size == 4 || size == 16) {
 		panic(fmt.Errorf("invalid IP size: %d", size))
 	}
 
-	if len(f.rbuf) < 1 {
-		panic(fmt.Errorf("not enough bytes in buffer to read inet require %d got: %d", size, len(f.rbuf)))
+	if len(f.buf) < 1 {
+		panic(fmt.Errorf("not enough bytes in buffer to read inet require %d got: %d", size, len(f.buf)))
 	}
 
 	ip := make([]byte, size)
-	copy(ip, f.rbuf[:size])
-	f.rbuf = f.rbuf[size:]
+	copy(ip, f.buf[:size])
+	f.buf = f.buf[size:]
 	return net.IP(ip)
 }
 
@@ -1923,7 +1914,7 @@ func (f *framer) readStringMultiMap() map[string][]string {
 }
 
 func (f *framer) writeByte(b byte) {
-	f.wbuf = append(f.wbuf, b)
+	f.buf = append(f.buf, b)
 }
 
 func appendBytes(p []byte, d []byte) []byte {
@@ -1980,29 +1971,29 @@ func (f *framer) writeCustomPayload(customPayload *map[string][]byte) {
 
 // these are protocol level binary types
 func (f *framer) writeInt(n int32) {
-	f.wbuf = appendInt(f.wbuf, n)
+	f.buf = appendInt(f.buf, n)
 }
 
 func (f *framer) writeUint(n uint32) {
-	f.wbuf = appendUint(f.wbuf, n)
+	f.buf = appendUint(f.buf, n)
 }
 
 func (f *framer) writeShort(n uint16) {
-	f.wbuf = appendShort(f.wbuf, n)
+	f.buf = appendShort(f.buf, n)
 }
 
 func (f *framer) writeLong(n int64) {
-	f.wbuf = appendLong(f.wbuf, n)
+	f.buf = appendLong(f.buf, n)
 }
 
 func (f *framer) writeString(s string) {
 	f.writeShort(uint16(len(s)))
-	f.wbuf = append(f.wbuf, s...)
+	f.buf = append(f.buf, s...)
 }
 
 func (f *framer) writeLongString(s string) {
 	f.writeInt(int32(len(s)))
-	f.wbuf = append(f.wbuf, s...)
+	f.buf = append(f.buf, s...)
 }
 
 func (f *framer) writeStringList(l []string) {
@@ -2028,13 +2019,13 @@ func (f *framer) writeBytes(p []byte) {
 		f.writeInt(-1)
 	} else {
 		f.writeInt(int32(len(p)))
-		f.wbuf = append(f.wbuf, p...)
+		f.buf = append(f.buf, p...)
 	}
 }
 
 func (f *framer) writeShortBytes(p []byte) {
 	f.writeShort(uint16(len(p)))
-	f.wbuf = append(f.wbuf, p...)
+	f.buf = append(f.buf, p...)
 }
 
 func (f *framer) writeConsistency(cons Consistency) {
diff --git a/session.go b/session.go
index 6d13209..3afdc6e 100644
--- a/session.go
+++ b/session.go
@@ -41,6 +41,7 @@ type Session struct {
 	batchObserver       BatchObserver
 	connectObserver     ConnectObserver
 	frameObserver       FrameHeaderObserver
+	streamObserver      StreamObserver
 	hostSource          *ringDescriber
 	stmtsLRU            *preparedLRU
 
@@ -161,6 +162,7 @@ func NewSession(cfg ClusterConfig) (*Session, error) {
 	s.batchObserver = cfg.BatchObserver
 	s.connectObserver = cfg.ConnectObserver
 	s.frameObserver = cfg.FrameHeaderObserver
+	s.streamObserver = cfg.StreamObserver
 
 	//Check the TLS Config before trying to connect to anything external
 	connCfg, err := connConfig(&s.cfg)


Test patch:
diff --git a/cassandra_test.go b/cassandra_test.go
index 77db0ed..8dc4a1f 100644
--- a/cassandra_test.go
+++ b/cassandra_test.go
@@ -2731,7 +2731,7 @@ func TestNegativeStream(t *testing.T) {
 	const stream = -50
 	writer := frameWriterFunc(func(f *framer, streamID int) error {
 		f.writeHeader(0, opOptions, stream)
-		return f.finishWrite()
+		return f.finish()
 	})
 
 	frame, err := conn.exec(context.Background(), writer, nil)
diff --git a/conn_test.go b/conn_test.go
index 2f836ae..845ce52 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -1,6 +1,7 @@
 // Copyright (c) 2012 The gocql Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
+//go:build all || unit
 // +build all unit
 
 package gocql
@@ -668,11 +669,14 @@ func TestStream0(t *testing.T) {
 	const expErr = "gocql: received unexpected frame on stream 0"
 
 	var buf bytes.Buffer
-	f := newFramer(nil, &buf, nil, protoVersion4)
+	f := newFramer(nil, protoVersion4)
 	f.writeHeader(0, opResult, 0)
 	f.writeInt(resultKindVoid)
-	f.wbuf[0] |= 0x80
-	if err := f.finishWrite(); err != nil {
+	f.buf[0] |= 0x80
+	if err := f.finish(); err != nil {
+		t.Fatal(err)
+	}
+	if err := f.writeTo(&buf); err != nil {
 		t.Fatal(err)
 	}
 
@@ -833,43 +837,44 @@ func TestWriteCoalescing(t *testing.T) {
 			t.Errorf("unexpected read error: %v", err)
 		}
 	}()
+	enqueued := make(chan struct{})
+	resetTimer := make(chan struct{})
 	w := &writeCoalescer{
+		writeCh: make(chan writeRequest),
 		c:       client,
-		writeCh: make(chan struct{}),
-		cond:    sync.NewCond(&sync.Mutex{}),
 		quit:    ctx.Done(),
-		running: true,
+		timeout: 500 * time.Millisecond,
+		testEnqueuedHook: func() {
+			enqueued <- struct{}{}
+		},
+		testFlushedHook: func() {
+			client.Close()
+		},
 	}
+	timerC := make(chan time.Time, 1)
+	go func() {
+		w.writeFlusherImpl(timerC, func() { resetTimer <- struct{}{} })
+	}()
 
 	go func() {
-		if _, err := w.Write([]byte("one")); err != nil {
+		if _, err := w.writeContext(context.Background(), []byte("one")); err != nil {
 			t.Error(err)
 		}
 	}()
 
 	go func() {
-		if _, err := w.Write([]byte("two")); err != nil {
+		if _, err := w.writeContext(context.Background(), []byte("two")); err != nil {
 			t.Error(err)
 		}
 	}()
 
-	bufMutex.Lock()
-	if buf.Len() != 0 {
-		t.Fatalf("expected buffer to be empty have: %v", buf.String())
-	}
-	bufMutex.Unlock()
+	<-enqueued
+	<-resetTimer
+	<-enqueued
 
-	for true {
-		w.cond.L.Lock()
-		if len(w.buffers) == 2 {
-			w.cond.L.Unlock()
-			break
-		}
-		w.cond.L.Unlock()
-	}
+	// flush
+	timerC <- time.Now()
 
-	w.flush()
-	client.Close()
 	<-done
 
 	if got := buf.String(); got != "onetwo" && got != "twoone" {
@@ -897,7 +902,7 @@ func TestWriteCoalescing_WriteAfterClose(t *testing.T) {
 	w := newWriteCoalescer(client, 0, 5*time.Millisecond, ctx.Done())
 
 	// ensure 1 write works
-	if _, err := w.Write([]byte("one")); err != nil {
+	if _, err := w.writeContext(context.Background(), []byte("one")); err != nil {
 		t.Fatal(err)
 	}
 
@@ -911,7 +916,7 @@ func TestWriteCoalescing_WriteAfterClose(t *testing.T) {
 	cancel()
 	client.Close() // close client conn too, since server won't see the answer anyway.
 
-	if _, err := w.Write([]byte("two")); err == nil {
+	if _, err := w.writeContext(context.Background(), []byte("two")); err == nil {
 		t.Fatal("expected to get error for write after closing")
 	} else if err != io.EOF {
 		t.Fatalf("expected to get EOF got %v", err)
@@ -1115,7 +1120,7 @@ func (srv *TestServer) serve() {
 					srv.onRecv(framer)
 				}
 
-				go srv.process(framer)
+				go srv.process(conn, framer)
 			}
 		}(conn)
 	}
@@ -1153,12 +1158,13 @@ func (srv *TestServer) errorLocked(err interface{}) {
 	srv.t.Error(err)
 }
 
-func (srv *TestServer) process(f *framer) {
-	head := f.header
+func (srv *TestServer) process(conn net.Conn, reqFrame *framer) {
+	head := reqFrame.header
 	if head == nil {
 		srv.errorLocked("process frame with a nil header")
 		return
 	}
+	respFrame := newFramer(nil, reqFrame.proto)
 
 	switch head.op {
 	case opStartup:
@@ -1170,12 +1176,12 @@ func (srv *TestServer) process(f *framer) {
 				return
 			}
 		}
-		f.writeHeader(0, opReady, head.stream)
+		respFrame.writeHeader(0, opReady, head.stream)
 	case opOptions:
-		f.writeHeader(0, opSupported, head.stream)
-		f.writeShort(0)
+		respFrame.writeHeader(0, opSupported, head.stream)
+		respFrame.writeShort(0)
 	case opQuery:
-		query := f.readLongString()
+		query := reqFrame.readLongString()
 		first := query
 		if n := strings.Index(query, " "); n > 0 {
 			first = first[:n]
@@ -1183,60 +1189,65 @@ func (srv *TestServer) process(f *framer) {
 		switch strings.ToLower(first) {
 		case "kill":
 			atomic.AddInt64(&srv.nKillReq, 1)
-			f.writeHeader(0, opError, head.stream)
-			f.writeInt(0x1001)
-			f.writeString("query killed")
+			respFrame.writeHeader(0, opError, head.stream)
+			respFrame.writeInt(0x1001)
+			respFrame.writeString("query killed")
 		case "use":
-			f.writeInt(resultKindKeyspace)
-			f.writeString(strings.TrimSpace(query[3:]))
+			respFrame.writeInt(resultKindKeyspace)
+			respFrame.writeString(strings.TrimSpace(query[3:]))
 		case "void":
-			f.writeHeader(0, opResult, head.stream)
-			f.writeInt(resultKindVoid)
+			respFrame.writeHeader(0, opResult, head.stream)
+			respFrame.writeInt(resultKindVoid)
 		case "timeout":
 			<-srv.ctx.Done()
 			return
 		case "slow":
 			go func() {
-				f.writeHeader(0, opResult, head.stream)
-				f.writeInt(resultKindVoid)
-				f.wbuf[0] = srv.protocol | 0x80
+				respFrame.writeHeader(0, opResult, head.stream)
+				respFrame.writeInt(resultKindVoid)
+				respFrame.buf[0] = srv.protocol | 0x80
 				select {
 				case <-srv.ctx.Done():
 					return
 				case <-time.After(50 * time.Millisecond):
-					f.finishWrite()
+					respFrame.finish()
+					respFrame.writeTo(conn)
 				}
 			}()
 			return
 		case "speculative":
 			atomic.AddInt64(&srv.nKillReq, 1)
 			if atomic.LoadInt64(&srv.nKillReq) > 3 {
-				f.writeHeader(0, opResult, head.stream)
-				f.writeInt(resultKindVoid)
-				f.writeString("speculative query success on the node " + srv.Address)
+				respFrame.writeHeader(0, opResult, head.stream)
+				respFrame.writeInt(resultKindVoid)
+				respFrame.writeString("speculative query success on the node " + srv.Address)
 			} else {
-				f.writeHeader(0, opError, head.stream)
-				f.writeInt(0x1001)
-				f.writeString("speculative error")
+				respFrame.writeHeader(0, opError, head.stream)
+				respFrame.writeInt(0x1001)
+				respFrame.writeString("speculative error")
 				rand.Seed(time.Now().UnixNano())
 				<-time.After(time.Millisecond * 120)
 			}
 		default:
-			f.writeHeader(0, opResult, head.stream)
-			f.writeInt(resultKindVoid)
+			respFrame.writeHeader(0, opResult, head.stream)
+			respFrame.writeInt(resultKindVoid)
 		}
 	case opError:
-		f.writeHeader(0, opError, head.stream)
-		f.wbuf = append(f.wbuf, f.rbuf...)
+		respFrame.writeHeader(0, opError, head.stream)
+		respFrame.buf = append(respFrame.buf, reqFrame.buf...)
 	default:
-		f.writeHeader(0, opError, head.stream)
-		f.writeInt(0)
-		f.writeString("not supported")
+		respFrame.writeHeader(0, opError, head.stream)
+		respFrame.writeInt(0)
+		respFrame.writeString("not supported")
 	}
 
-	f.wbuf[0] = srv.protocol | 0x80
+	respFrame.buf[0] = srv.protocol | 0x80
+
+	if err := respFrame.finish(); err != nil {
+		srv.errorLocked(err)
+	}
 
-	if err := f.finishWrite(); err != nil {
+	if err := respFrame.writeTo(conn); err != nil {
 		srv.errorLocked(err)
 	}
 }
@@ -1247,9 +1258,9 @@ func (srv *TestServer) readFrame(conn net.Conn) (*framer, error) {
 	if err != nil {
 		return nil, err
 	}
-	framer := newFramer(conn, conn, nil, srv.protocol)
+	framer := newFramer(nil, srv.protocol)
 
-	err = framer.readFrame(&head)
+	err = framer.readFrame(conn, &head)
 	if err != nil {
 		return nil, err
 	}
diff --git a/frame_test.go b/frame_test.go
index b846e60..6b8eb22 100644
--- a/frame_test.go
+++ b/frame_test.go
@@ -35,16 +35,14 @@ func TestFuzzBugs(t *testing.T) {
 	for i, test := range tests {
 		t.Logf("test %d input: %q", i, test)
 
-		var bw bytes.Buffer
-
 		r := bytes.NewReader(test)
 		head, err := readHeader(r, make([]byte, 9))
 		if err != nil {
 			continue
 		}
 
-		framer := newFramer(r, &bw, nil, byte(head.version))
-		err = framer.readFrame(&head)
+		framer := newFramer(nil, byte(head.version))
+		err = framer.readFrame(r, &head)
 		if err != nil {
 			continue
 		}
@@ -64,12 +62,11 @@ func TestFrameWriteTooLong(t *testing.T) {
 		t.Skip("skipping test in travis due to memory pressure with the race detecor")
 	}
 
-	w := &bytes.Buffer{}
-	framer := newFramer(nil, w, nil, 2)
+	framer := newFramer(nil, 2)
 
 	framer.writeHeader(0, opStartup, 1)
 	framer.writeBytes(make([]byte, maxFrameSize+1))
-	err := framer.finishWrite()
+	err := framer.finish()
 	if err != ErrFrameTooBig {
 		t.Fatalf("expected to get %v got %v", ErrFrameTooBig, err)
 	}
@@ -85,7 +82,7 @@ func TestFrameReadTooLong(t *testing.T) {
 	// write a new header right after this frame to verify that we can read it
 	r.Write([]byte{0x02, 0x00, 0x00, byte(opReady), 0x00, 0x00, 0x00, 0x00})
 
-	framer := newFramer(r, nil, nil, 2)
+	framer := newFramer(nil, 2)
 
 	head := frameHeader{
 		version: 2,
@@ -93,7 +90,7 @@ func TestFrameReadTooLong(t *testing.T) {
 		length:  r.Len() - 8,
 	}
 
-	err := framer.readFrame(&head)
+	err := framer.readFrame(r, &head)
 	if err != ErrFrameTooBig {
 		t.Fatalf("expected to get %v got %v", ErrFrameTooBig, err)
 	}
diff --git a/framer_bench_test.go b/framer_bench_test.go
index 06dfac4..5b7b442 100644
--- a/framer_bench_test.go
+++ b/framer_bench_test.go
@@ -37,7 +37,7 @@ func BenchmarkParseRowsFrame(b *testing.B) {
 				op:      opResult,
 				length:  len(data),
 			},
-			rbuf: data,
+			buf: data,
 		}
 
 		_, err = framer.parseFrame()