Scale Stack 0.3 Developer Documentation

scalestack/database_proxy/connection.cc

Go to the documentation of this file.

/*
 * Scale Stack
 *
 * Copyright 2010 Eric Day
 *
 * 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.
 */

#include "config.h"

#include <cstring>

#include <scalestack/database_proxy/service.h>
#include <scalestack/database_proxy/client.h>
#include <scalestack/database_proxy/connection.h>
#include <scalestack/kernel/module.h>

namespace scalestack
{
namespace database_proxy
{

/*
 * Public methods.
 */

connection::connection(service& service):
  network::stream(service),
  _state(),
  _need_handshake(true),
  _need_yield(false),
  _in_transaction(false),
  _shutdown(false),
  _size(),
  _packet_size(),
  _buffer(),
  _client()
{
  _client = service.get_client(this, _need_handshake);
}

size_t connection::read(uint8_t* buffer, size_t size)
{
  _size = size;
  _buffer = buffer;
  return _client_write();
}

void connection::continue_write(void)
{
  consume(_client_write());
}

void connection::flush_write(void)
{
  if (_client != NULL)
    _client->continue_write();
}

void connection::error(void)
{
  _shutdown = true;
  run_now();
}

void connection::run(void)
{
  if (_shutdown)
  {
    if (_client == NULL)
    {
      static_cast<service&>(_stream_service).remove_from_queue(this);
      stop();
      return;
    }

    static_cast<service&>(_stream_service).yield_client(_client);
    _client = NULL;
    return;
  }

  if (_client == NULL)
    _client = static_cast<service&>(_stream_service).get_client(this, _need_handshake);
  else if (_need_yield)
  {
    _need_yield = false;
    static_cast<service&>(_stream_service).yield_client(_client);
    _client = NULL;
  }
}

void connection::yield(void)
{
  _need_yield = true;
  run_now();
}

/*
 * Private methods.
 */

size_t connection::_client_write(void)
{
  if (_shutdown)
    return 0;

  if (_need_yield)
  {
    _need_yield = false;
    static_cast<service&>(_stream_service).yield_client(_client);
    _client = NULL;
  }

  if (_client == NULL)
  {
    _client = static_cast<service&>(_stream_service).get_client(this, _need_handshake);
    /* Just return here, the client will call this when ready. */
    return 0;
  }

  size_t written = 0;

  switch (_state)
  {
  case CLIENT_HANDSHAKE_PACKET:
    if (_size < 4)
      return 0;

    _packet_size = 4 + (_buffer[0] | (static_cast<size_t>(_buffer[1]) << 8) | (static_cast<size_t>(_buffer[2]) << 16));
    _state = CLIENT_HANDSHAKE_FLUSH;
    _module.log_debug("C-HAND packet: %u", _packet_size);

  case CLIENT_HANDSHAKE_FLUSH:
    written = _client->write(_buffer, _packet_size, true);
    _size -= written;
    _buffer += written;

    _packet_size -= written;
    if (_packet_size == 0)
    {
      _state = COMMAND_PACKET;
      _need_handshake = false;
    }

    break;

  case COMMAND_PACKET:
    if (_size < 4)
      return 0;

    _packet_size = 4 + (_buffer[0] | (static_cast<size_t>(_buffer[1]) << 8) | (static_cast<size_t>(_buffer[2]) << 16));
    _state = COMMAND_FLUSH;
    _module.log_debug("COMMAND packet: %u", _packet_size);
    if (_packet_size == 10 && _size >= 10 && !memcmp("BEGIN", _buffer + 5, 5))
      _in_transaction = true;
    else if (!_in_transaction ||
             (_packet_size == 11 && _size >= 11 &&
              !memcmp("COMMIT", _buffer + 5, 6)))
    {
      _in_transaction = false;
    }

    if (_packet_size == 5 && _buffer[4] == 1)
    {
      error();
      return 5;
    }

  case COMMAND_FLUSH:
    written = _client->write(_buffer, _packet_size, true);
    _size -= written;
    _buffer += written;

    _packet_size -= written;
    if (_packet_size == 0)
      _state = COMMAND_PACKET;

    break;
  }

  if (!_in_transaction && _packet_size == 0)
    _client->yield();

  return written;
}

} /* namespace database_proxy */
} /* namespace scalestack */

---