source: gs2-extensions/tdb-edit/trunk/src/perllib/dbutil/tdbserver.pm@ 28765

###########################################################################
#
# dbutil::tdbserver -- utility functions for writing to tdb databases but
#                      implemented as a server with a single, persistent
#                      connection
#
# A component of the Greenstone digital library software
# from the New Zealand Digital Library Project at the
# University of Waikato, New Zealand.
#
# Copyright (C) 2012
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
#
###########################################################################

package dbutil::tdbserver;

# Pragma
use strict;
use warnings;

# Modules
# - we're going to have to delve into locking (a little) to prevent
#   multiple threads trying to launch the server at once
use Fcntl qw(:flock);
# - we'll also need to wait around for an indeterminate amount of time for the
#   TDBServer to exit (detected by its lockfile disappearing)
use Time::HiRes  qw( usleep );

# Greenstone modules
use TDBClient;
use util;

my $hyphen70 = '-' x 70;
my $debug = 0;

# We have a global reference to all of the TDB Server lockfiles that this
# instance has created (as we'll be responsible for closing them)
my %created_server_lockfile_paths;
# Keep track of the lockfiles for server we have added ourselves as listeners
# to.
my %listener_server_lockfile_paths;
# We also have a global of all of the listeners we have assigned as we'll
# be responsible for removing them.
my %registered_listeners;

sub _spawnClient
{
  my ($infodb_file_path) = @_;
  # 1. Check whether the server is already running by trying to locate the
  #    server 'lock' file.

  my ($collection) = $infodb_file_path =~ /collect[\\\/]([^\\\/]+)/i;
  my ($infodb_file, $extension) = $infodb_file_path =~ /([^\\\/]+)\.(t?db)/i;

  my $tmp_dir = &util::filename_cat($ENV{'GSDLHOME'},'tmp');
  if (!-d $tmp_dir)
  {
    mkdir($tmp_dir, 0755);
  }

  my $we_started_server = 0;
  my $server_lockfile_path =  &util::filename_cat($ENV{'GSDLHOME'},'collect',$collection,'tmp','tdbserver.lock');
  print " * Searching for TDBServer lockfile...\n" if ($debug);
  if (!-e $server_lockfile_path)
  {
    print "Not found.\n" if ($debug);
    # We need to lock here to ensure only one thread enters the following code,
    # sees a missing TDBServer, and launches it
    my $tmp_lockfile_path = &util::filename_cat($ENV{'GSDLHOME'},'tmp','dbutil-tdbserver.lock');
    open(TMPFH, '>', $tmp_lockfile_path) or die("Warning! Failed to open file for writing: " . $tmp_lockfile_path . "\nReason: " . $! . "\n");
    flock(TMPFH, LOCK_EX) or die("Error! Cannot lock file exclusively: " . $tmp_lockfile_path . "\nReason: " . $! . "\n");
    print TMPFH localtime();
    # - If the file still doesn't exist...
    if (!-e $server_lockfile_path)
    {
      $we_started_server = 1;
      # ...start it!
      my $launch_cmd = 'TDBServer.pl "' . $$ . '" "' . $collection . '"';
      print "* Starting TDBServer for: " . $collection . " [" . $launch_cmd . "]...\n";
      # @note I once had the below pipe ending with 2>&1 |, but that then blocks
      #       indefinitely when looping and reading <SERVERIN>.
      open(SERVERIN, $launch_cmd . ' |') or die("Error! Failed to run launch command: " . $launch_cmd . "\nReason: " . $! . "\n");
      # read all the output from the server
      my $line = '';
      my $server_lock_file_created = 0;
      autoflush STDOUT 1;
      print "* Waiting for TDBServer to start";
      my $server_host = '';
      my $server_port = '';
      while ($line = <SERVERIN>)
      {
        print ".";
        # - watch for the line indicating a lock file has been created and
        #   populated with a sexy port number
        if ($line =~ /Server now listening on ([^:]+):(\d+)/)
        {
          $server_host = $1;
          $server_port = $2;
          $server_lock_file_created = 1;
        }
        # - we could also watch for errors here
        if ($debug)
        {
          if ($line !~ /\n/)
          {
            $line .= "\n";
          }
          $|++; # autoflush
          print "[tdbserver] $line";
          $|--; # disable autoflush
        }
      }
      close(SERVERIN);
      print "\n";
      if (!$server_lock_file_created)
      {
        die("Error! TDBServer failed to create lock file. Check server logs.");
      }
      elsif ($debug)
      {
        print "* Server now running on " . $server_host . ":" . $server_port . "\n";
      }
      # record this for later
      $created_server_lockfile_paths{$server_lockfile_path} = 1;
    }
    flock(TMPFH, LOCK_UN);
    close($tmp_lockfile_path);
    unlink($tmp_lockfile_path);
  }
  else
  {
    print "Found!\n" if ($debug);
  }
  # record this for later
  $listener_server_lockfile_paths{$server_lockfile_path} = $infodb_file_path;
  return TDBClient->new($server_lockfile_path, $infodb_file, $we_started_server);
}

END
{
  # we ask the server to shutdown, but only the 'creator' thread will actually
  # be able to, and only once all listeners have deregistered.
  foreach my $server_lockfile_path (keys (%listener_server_lockfile_paths))
  {
    my $infodb_file_path = $listener_server_lockfile_paths{$server_lockfile_path};
    my $tdb_client_handle = TDBClient->new($server_lockfile_path, '');
    # Deregister all of our registered listeners
    foreach my $listener_suffix (keys(%registered_listeners))
    {
      $tdb_client_handle->removeListener($listener_suffix);
    }
    # ask the servers we created to shut down (all other threads will have
    # this request ignored)
    if (defined $created_server_lockfile_paths{$infodb_file_path})
    {
      print "* Attempting to stop TDBServer for: " . $infodb_file_path . "\n";
    }
    $tdb_client_handle->stopServer();
  }
  # we should now wait until all of our server_lockfiles have actually been
  # removed (otherwise people could mistakenly run import/build again
  # immediately and things *might* go pearshaped). (Subsequent testing shows it
  # won't - a new Greenstone build can't start until the previous lockfile is
  # removed - but I think, for ATOMic sake, I may as well make waiting a part
  # of the Greenstone import/build time.
  my $wait_for_exit = 1;
  foreach my $server_lockfile_path (keys (%created_server_lockfile_paths))
  {
    if ($wait_for_exit)
    {
      # While the file exists, we should wait
      autoflush STDOUT 1;
      my $blurb = '';
      if ($debug)
      {
        $blurb = '[' . $server_lockfile_path . ']';
      }
      print '* Waiting for TDBServer to exit... ';
      if (-e $server_lockfile_path)
      {
        while (-e $server_lockfile_path)
        {
          usleep(10000);
        }
      }
    }
    else
    {
      print '* NOT Waiting for TDBServer to exit... ';
    }
    print " Done!\n";
  }
}

# -----------------------------------------------------------------------------
#   TDB SERVER IMPLEMENTATION
# -----------------------------------------------------------------------------
sub open_infodb_write_handle
{
  my $infodb_file_path = shift(@_);
  my $opt_append = shift(@_);
  if (defined $opt_append && $opt_append ne "append")
  {
    print "Warning! Modes other than 'append' not supported for TDBServer.\n";
  }
  my $tdb_client_handle = &_spawnClient($infodb_file_path);
  # Register this client on the server if necessary
  $tdb_client_handle->addListener('w');
  $registered_listeners{'w'} = 1;
  # and pass the handle to client around
  return $tdb_client_handle;
}

# /** Destructor or near enough.
#  /*
sub close_infodb_write_handle
{
  my $tdb_client_handle = shift(@_);
  # @todo Is there meant to be something here?
}
# /** close_infodb_write_handle($infodb_handle) **/

# /** @function get_info_db_file_path
#  *  Exactly the same as vanilla TDB - as we are still using a TDB database
#  *  just accessing it via a persistant server
#  */
sub get_infodb_file_path
{
  my $collection_name = shift(@_);
  my $infodb_directory_path = shift(@_);
  my $create_server = shift(@_);

  my $infodb_file_extension = ".tdb";
  my $infodb_file_name = &util::get_dirsep_tail($collection_name) . $infodb_file_extension;
  my $infodb_file_path = &util::filename_cat($infodb_directory_path, $infodb_file_name);

  # Special Case for TDBServer
  if (defined $create_server && $create_server == 1)
  {
    my $tdb_client_handle = &_spawnClient($infodb_file_path);
    # Register this client on the server if necessary
    $tdb_client_handle->addListener('i');
    $registered_listeners{'i'} = 1;
  }

  # Resuming our regular programming
  return $infodb_file_path;
}

sub read_infodb_file
{
  my $infodb_file_path = shift(@_);
  my $infodb_map = shift(@_);

  my $tdb_client_handle = &_spawnClient($infodb_file_path);
  $tdb_client_handle->addListener('r');
  $registered_listeners{'r'} = 1;
  # retrieves all the keys in the form:
  # [key1]\n[key2]\n[key3]\n...[keyn]
  my $raw_infodb_keys = $tdb_client_handle->query('[*]?');

  my @infodb_keys = split(/\r?\n/, $raw_infodb_keys);
  foreach my $infodb_key (@infodb_keys)
  {
    if ($infodb_key =~ /.+/ && $infodb_key !~ /-{70}/)
    {
      # lookup each key
      my $infodb_value = $tdb_client_handle->query('[' . $infodb_key . ']?');
      # store it
      $infodb_map->{$infodb_key} = $infodb_value;
    }
  }
}

sub read_infodb_keys
{
  my $infodb_file_path = shift(@_);
  my $infodb_map = shift(@_);

  # spawn client (creating server as necessary)
  my $tdb_client_handle = &_spawnClient($infodb_file_path);
  # register ourself as listener
  $tdb_client_handle->addListener('k');
  $registered_listeners{'k'} = 1;
  # retrieves all the keys in the form:
  # [key1]\n[key2]\n[key3]\n...[keyn]
  my $raw_infodb_keys = $tdb_client_handle->query('[*]?');
  my @infodb_keys = split(/\r?\n/, $raw_infodb_keys);
  foreach my $infodb_key (@infodb_keys)
  {
    if ($infodb_key =~ /.+/ && $infodb_key !~ /-{70}/)
    {
      $infodb_map->{$infodb_key} = 1;
    }
  }
}

sub write_infodb_entry
{
  my $tdb_client_handle = shift(@_);
  my $infodb_key = shift(@_);
  my $infodb_map = shift(@_);
  # - build up the tdb command
  my $tdb_command = "[" . $infodb_key . "]+\n";
  foreach my $infodb_value_key (keys(%$infodb_map))
  {
    foreach my $infodb_value (@{$infodb_map->{$infodb_value_key}})
    {
      if ($infodb_value =~ /-{70,}/)
      {
        # if value contains 70 or more hyphens in a row we need to escape them
        # to prevent txt2db from treating them as a separator
        $infodb_value =~ s/-/&\#045;/gi;
      }
      $tdb_command .= "<" . $infodb_value_key . ">" . $infodb_value . "\n";
    }
  }
  $tdb_command .= $hyphen70 . "\n";
  # - ask the client to transmit the command to the server
  $tdb_client_handle->query($tdb_command);
}

sub write_infodb_rawentry
{
  my $tdb_client_handle = shift(@_);
  my $infodb_key = shift(@_);
  my $infodb_val = shift(@_);
  # - build up the tdb command
  my $tdb_command = "[" . $infodb_key . "]\n";
  $tdb_command .= $infodb_val . "\n";
  $tdb_command .= $hyphen70 . "\n";
  # - ask the client to transmit the command to the server
  $tdb_client_handle->query($tdb_command);
  return 1;
}

sub set_infodb_entry
{
  my $infodb_file_path = shift(@_);
  my $infodb_key = shift(@_);
  my $infodb_map = shift(@_);
  # spawn client (creating server as necessary)
  my $tdb_client_handle = &_spawnClient($infodb_file_path);
  $tdb_client_handle->addListener('s');
  $registered_listeners{'s'} = 1;
  # Protect metadata values that go inside quotes for tdbset
  foreach my $k (keys %$infodb_map)
  {
    my @escaped_v = ();
    foreach my $v (@{$infodb_map->{$k}})
    {
      if ($k eq "contains")
      {
        # protect quotes in ".2;".3 etc
        $v =~ s/\"/\\\"/g;
        push(@escaped_v, $v);
      }
      else
      {
        my $ev = &ghtml::unescape_html($v);
        $ev =~ s/\"/\\\"/g;
        push(@escaped_v, $ev);
      }
    }
    $infodb_map->{$k} = \@escaped_v;
  }
  # Generate the record string (TDB command)
  my $tdb_command = "[" . $infodb_key . "]\n";
  $tdb_command .= &dbutil::convert_infodb_hash_to_string($infodb_map) . "\n";
  $tdb_command .= $hyphen70 . "\n";
  # Send command to server
  $tdb_client_handle->query($tdb_command);
}

sub delete_infodb_entry
{
  my $tdb_client_handle = shift(@_);
  my $infodb_key = shift(@_);
  # - create command
  my $tdb_command = "[" . $infodb_key . "]-\n";
  # - and send
  $tdb_client_handle->query($tdb_command);
}



1;